Unnecessary check for trace level in nnx.Variable.__setattr__
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
In the development of a quantization library, we often need to collect some statistics of activations. Sometimes, the collection happens inside a custom_vjp function, as demonstrated below.
```py
class QuantStats(nnx.Variable):
# __setattr__ = object.__setattr__
pass
class Model(nnx.Module):
def __init__(self):
self.stats = QuantStats({'absmax': jnp.zeros(())})
self.linear = nnx.Linear(12, 10, rngs=nnx.Rngs(0))
def __call__(self, x):
@jax.custom_vjp
def f(x):
return fwd(x)[0]
def fwd(x):
self.stats.value = {'absmax': jnp.max(jnp.abs(x))}
return x, ()
def bwd(_, g):
return g
f.defvjp(fwd, bwd)
return self.linear(f(x))
def loss_fn(model, x):
out = model(x)
return jnp.sum(jnp.abs(out))
model = Model()
loss_fn(model, jnp.full((1, 12), 42.))
print(model.stats['absmax'])
nnx.grad(loss_fn)(model, jnp.full((1, 12), 43.))
print(model.stats['absmax'])
```
Today, running the above code will raise an error like this
```
/tmp/ipython-input-90-2380234258.py in fwd(x)
15
16 def fwd(x):
---> 17 self.stats.value = {'absmax': jnp.max(jnp.abs(x))}
18 return x, ()
19
.../flax/nnx/variablelib.py in __setattr__(self, name, value)
274 name != 'value' or not self.mutable
275 ):
--> 276 raise errors.TraceContextError(
277 f'Cannot mutate {type(self).__name__} from a different trace level'
278 )
TraceContextError: Cannot mutate QuantStats from a different trace level (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.TraceContextError)
```
However, jax doesn't require a custom_vjp function to be pure. The above code will work if we uncomment the `__setattr__ = object.__setattr__` line.
It seems that NNX is imposing an unnecessary check here.
Contributor guide
Assessment
This issue has not been assessed yet.