Bad interaction between nnx.Rngs and custom derivatives
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
```
File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.11/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 177, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^
File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.11/lib/python3.11/site-packages/jax/_src/custom_derivatives.py", line 613, in __call__
in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.11/lib/python3.11/site-packages/jax/_src/custom_derivatives.py", line 613, in
in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat]
^^^^^^^^^^^^^^^^
File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.11/lib/python3.11/site-packages/jax/_src/core.py", line 1431, in get_aval
return concrete_aval(x)
^^^^^^^^^^^^^^^^
File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.11/lib/python3.11/site-packages/jax/_src/core.py", line 1422, in concrete_aval
return concrete_aval(x.__jax_array__())
^^^^^^^^^^^^^^^^^
File "/home/neil/src/flax/flax/experimental/nnx/nnx/rnglib.py", line 148, in
return lambda: self._make_rng(name)
^^^^^^^^^^^^^^^^^^^^
File "/home/neil/src/flax/flax/experimental/nnx/nnx/rnglib.py", line 140, in _make_rng
raise ValueError(f"No RNG named {name!r} or 'default' found in Rngs.")
ValueError: No RNG named '__jax_array__' or 'default' found in Rngs.
```
Contributor guide
Assessment
This issue has not been assessed yet.