Comments (5)
Hmmm I see. This is a fix, by making a wrapper object basically for the tuple singleton: #717. This is definitely the craziest edge case to write this much code for lol
from equinox.
Not 100% sure on the source, but one fix is to force it to recognize tuples as endpoints, via
import equinox as eqx
class M(eqx.Module):
a: tuple
b: tuple
def __init__(self):
self.a = ()
self.b = (1,)
eqx.tree_at(lambda m: m.a, M(), (), is_leaf=lambda x: x == ())
from equinox.
Interesting! Looks like Python actually uses the exact same object for tuple literals:
() is () # True
(1,) is (1,) # True
So the logic inside eqx.tree_at
works by requiring that for any two nodes x, y
in the PyTree, that x is not y
. This is how it can identify what to replace. (Here I use "node" to refer to both composite nodes and leaves.)
Now specifically leaves frequently fail this (e.g. x = [4, 4]; x[0] is x[1] # True
), so we make this work by first doing a pytree = jtu.tree_map(Wrapper, pytree)
to wrap up the leaves. No big deal.
And until this point, there haven't been any examples of composite nodes which fail the x is not y
check. Now it seems we have one! This is why two tuples are specifically what is required to bump into this issue.
If anyone feels like tackling this in the near future then I'd be happy to take a pull request.
from equinox.
Is this something that needs a fix, it seems like its working as it should? This exists because an empty tuple is a singleton in CPython, so only 1 ever exists in memory (so all () have the same id) (https://stackoverflow.com/questions/14135542/how-is-tuple-implemented-in-cpython/). If you have something like
class M(eqx.Module):
a: tuple
b: tuple
def __init__(self):
self.a = (1,)
self.b = (1,)
this works because the id's are different. In my mind, equinox is behaving as it should, the where function is underspecified since it points to two objects in memory and its hard to think of a workaround that works for making tuples which are identical in memory that doesn't break other things that are equal in memory (outside of a hardcoded if x == ()
or something). Even if there was a workaround, I think it wouldn't be ideal for two reasons 1) it's such very much an edge case, 2) people who knew about tuples being singletons going into equinox (I just learned it) would then be surprised by the behavior because now equinox is going against what CPython does.
from equinox.
So I think of eqx.tree_at
as just being a "please change the node at this path in the PyTree". The fact that it goes via CPython's id
is an implementation detail, not its defining characteristic.
I think on that basis this is a minor bug. Though as you say, this is a very very edge case with a loud failure and an easy work-around, so it's not one I'm too worried by overall!
from equinox.
Related Issues (20)
- Freezing layers with `eqx.tree_map` (not parameters in a single layer) HOT 2
- Realistic scenario in which jax array is static HOT 4
- eqx.filter_hessian gives incorrect result for jax.nn.relu HOT 1
- [Question] Modifying the Static Variable in a Model HOT 5
- [Question/Maybe Issue] Is there a version of Jax.linearize in equinox HOT 3
- Add Tensor class HOT 3
- Adding a kernel constraint HOT 2
- Spectral Norm Features HOT 1
- Non-default `__init__` with `default_factory` throws error. HOT 3
- Comparable PixelShuffle functionality? HOT 2
- Question about efficiency of higher-order derivatives when defining custom jvp HOT 1
- Type errors with Module HOT 2
- PositionalSharding tensor of different ranks: use sharding.replicate! HOT 3
- Minimal Transformer encoder using MultiheadAttention and RoPE HOT 3
- debug.assert_max_traces is not helpful when the difference is due to weak_type HOT 1
- _DistinctTuple is not a tuple causing errors in tree_at HOT 1
- MyPy doesn't understand Module HOT 2
- Spliting a module in more than 2 partitions HOT 7
- Native Model Parallellism in equinox HOT 9
- v0.4.11.5 Breaks Pycharm Debugger Console HOT 11
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from equinox.