google / google/flax

Bad interaction between nnx.Rngs and custom derivatives

Open
#3,528 6 comments 0 reactions 1 assignee Claimed by @cgarciae View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.