Allow __setitem__ on tracers? #8519
Replies: 1 comment 4 replies
-
This is an interesting thought, but my first inclination is that we would not want to add this. First of all, one of the guiding design principles behind JAX is to strive for JIT invariance: that is, as much as possible code executed under JIT should behave the same as in op-by-op mode. This would break that assumption. Additionally, I think it would cause a lot of confusion because we could only make it work for So the idea of overriding only Others may have different reactions though - thanks for starting the discussion! |
Beta Was this translation helpful? Give feedback.
-
I am somewhat of a power user of JAX, but I still can't quite get used to JAX's syntax for setitem operations. To be specific, I am talking about this:
Concrete problems:
import numpy as npx
orimport jax.numpy as npx
).arr =
on the left hand side, so this silently turns into a no-op.Proposal:
I assume that JAX arrays being immutable is a non-negotiable requirement, but maybe that doesn't need to extend to Tracer objects? In this case, JAX could support
__setitem__
only under JIT:What do you think?
Beta Was this translation helpful? Give feedback.
All reactions