nnx.vmap example from documentation raise an index error
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
I am encountering an index error when running [this](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/nnx_basics.html#:~:text=from%20functools%20import%20partial) example in the documentation
I am running the code in a docker environment using an NVIDIA image for jax.
Best
Giovanni
### System information
- Docker Image ghcr.io/nvidia/jax:base
- flax==0.9.0
- jax[cuda12_local]==0.4.34
### Problem you have encountered:
Error while going through nnx tutorial
### What you expected to happen:
### Logs, error messages, etc:
---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
Cell In[8], [line 23](vscode-notebook-cell:?execution_count=8&line=23)
[19](vscode-notebook-cell:?execution_count=8&line=19) @partial(nnx.vmap, axis_size=5)
[20](vscode-notebook-cell:?execution_count=8&line=20) def create_model(rngs: nnx.Rngs):
[21](vscode-notebook-cell:?execution_count=8&line=21) return MLP(10, 32, 10, rngs=rngs)
---> [23](vscode-notebook-cell:?execution_count=8&line=23) model = create_model(nnx.Rngs(0))
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/graph.py:1158, in UpdateContextManager.__call__..update_context_manager_wrapper(*args, **kwargs)
[1155](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/graph.py:1155) @functools.wraps(f)
[1156](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/graph.py:1156) def update_context_manager_wrapper(*args, **kwargs):
[1157](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/graph.py:1157) with self:
-> [1158](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/graph.py:1158) return f(*args, **kwargs)
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:339, in vmap..vmap_wrapper(*args, **kwargs)
[335](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:335) args = resolve_kwargs(f, args, kwargs)
[336](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:336) pure_args = extract.to_tree(
[337](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:337) args, prefix=in_axes, split_fn=_vmap_split_fn, ctxtag='vmap'
[338](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:338) )
--> [339](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:339) pure_args_out, pure_out = vmapped_fn(*pure_args)
[340](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:340) _args_out, out = extract.from_tree((pure_args_out, pure_out), ctxtag='vmap')
[341](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:341) return out
[... skipping hidden 3 frame]
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:164, in VmapFn.__call__(self, *pure_args)
[159](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:159) pure_args = _update_variable_sharding_metadata(
[160](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:160) pure_args, self.transform_metadata, spmd.remove_axis
[161](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:161) )
[162](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:162) args = extract.from_tree(pure_args, ctxtag='vmap')
--> [164](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:164) out = self.f(*args)
[166](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:166) args_out = extract.clear_non_graph_nodes(args)
[167](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:167) pure_args_out, pure_out = extract.to_tree(
[168](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:168) (args_out, out),
[169](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:169) prefix=(self.in_axes, self.out_axes),
[170](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:170) split_fn=_vmap_split_fn,
[171](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:171) ctxtag='vmap',
[172](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/transforms/iteration.py:172) )
Cell In[8], [line 21](vscode-notebook-cell:?execution_count=8&line=21)
[19](vscode-notebook-cell:?execution_count=8&line=19) @partial(nnx.vmap, axis_size=5)
[20](vscode-notebook-cell:?execution_count=8&line=20) def create_model(rngs: nnx.Rngs):
---> [21](vscode-notebook-cell:?execution_count=8&line=21) return MLP(10, 32, 10, rngs=rngs)
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:79, in ObjectMeta.__call__(cls, *args, **kwargs)
[78](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:78) def __call__(cls, *args: Any, **kwargs: Any) -> Any:
---> [79](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:79) return _graph_node_meta_call(cls, *args, **kwargs)
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:88, in _graph_node_meta_call(cls, *args, **kwargs)
[86](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:86) node = cls.__new__(cls, *args, **kwargs)
[87](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:87) vars(node)['_object__state'] = ObjectState()
---> [88](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:88) cls._object_meta_construct(node, *args, **kwargs)
[90](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:90) return node
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:82, in ObjectMeta._object_meta_construct(cls, self, *args, **kwargs)
[81](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:81) def _object_meta_construct(cls, self, *args, **kwargs):
---> [82](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:82) self.__init__(*args, **kwargs)
Cell In[8], [line 8](vscode-notebook-cell:?execution_count=8&line=8)
[7](vscode-notebook-cell:?execution_count=8&line=7) def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):
----> [8](vscode-notebook-cell:?execution_count=8&line=8) self.linear1 = nnx.Linear(din, dmid, rngs=rngs)
[9](vscode-notebook-cell:?execution_count=8&line=9) self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
[10](vscode-notebook-cell:?execution_count=8&line=10) self.bn = nnx.BatchNorm(dmid, rngs=rngs)
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:79, in ObjectMeta.__call__(cls, *args, **kwargs)
[78](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:78) def __call__(cls, *args: Any, **kwargs: Any) -> Any:
---> [79](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:79) return _graph_node_meta_call(cls, *args, **kwargs)
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:88, in _graph_node_meta_call(cls, *args, **kwargs)
[86](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:86) node = cls.__new__(cls, *args, **kwargs)
[87](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:87) vars(node)['_object__state'] = ObjectState()
---> [88](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:88) cls._object_meta_construct(node, *args, **kwargs)
[90](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:90) return node
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:82, in ObjectMeta._object_meta_construct(cls, self, *args, **kwargs)
[81](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:81) def _object_meta_construct(cls, self, *args, **kwargs):
---> [82](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/object.py:82) self.__init__(*args, **kwargs)
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:346, in Linear.__init__(self, in_features, out_features, use_bias, dtype, param_dtype, precision, kernel_init, bias_init, dot_general, rngs)
[332](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:332) def __init__(
[333](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:333) self,
[334](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:334) in_features: int,
(...)
[344](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:344) rngs: rnglib.Rngs,
[345](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:345) ):
--> [346](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:346) kernel_key = rngs.params()
[347](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:347) self.kernel = nnx.Param(
[348](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:348) kernel_init(kernel_key, (in_features, out_features), param_dtype)
[349](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:349) )
[350](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/nn/linear.py:350) if use_bias:
File /usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/rnglib.py:84, in RngStream.__call__(self)
[80](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/rnglib.py:80) def __call__(self) -> jax.Array:
[81](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/rnglib.py:81) self.check_valid_context(
[82](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/rnglib.py:82) lambda: 'Cannot call RngStream from a different trace level'
[83](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/rnglib.py:83) )
---> [84](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/rnglib.py:84) key = jax.random.fold_in(self.key.value, self.count.value)
[85](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/rnglib.py:85) self.count.value += 1
[86](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/flax/nnx/nnx/rnglib.py:86) return key
File /usr/local/lib/python3.10/dist-packages/jax/_src/random.py:262, in fold_in(key, data)
[251](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:251) def fold_in(key: KeyArrayLike, data: IntegerArray) -> KeyArray:
[252](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:252) """Folds in data to a PRNG key to form a new PRNG key.
[253](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:253)
[254](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:254) Args:
(...)
[260](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:260) statistically safe for producing a stream of new pseudo-random values.
[261](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:261) """
--> [262](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:262) key, wrapped = _check_prng_key("fold_in", key)
[263](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:263) if np.ndim(data):
[264](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:264) raise TypeError("fold_in accepts a scalar, but was given an array of"
[265](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:265) f"shape {np.shape(data)} != (). Use jax.vmap for batching.")
File /usr/local/lib/python3.10/dist-packages/jax/_src/random.py:74, in _check_prng_key(name, key, allow_batched)
[72](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:72) def _check_prng_key(name: str, key: KeyArrayLike, *,
[73](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:73) allow_batched: bool = False) -> tuple[KeyArray, bool]:
---> [74](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:74) if isinstance(key, Array) and dtypes.issubdtype(key.dtype, dtypes.prng_key):
[75](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:75) wrapped_key = key
[76](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/random.py:76) wrapped = False
[... skipping hidden 1 frame]
File /usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/batching.py:346, in BatchTracer.aval(self)
[344](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/batching.py:344) return aval
[345](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/batching.py:345) elif type(self.batch_dim) is int:
--> [346](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/batching.py:346) return core.mapped_aval(aval.shape[self.batch_dim], self.batch_dim, aval)
[347](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/batching.py:347) elif type(self.batch_dim) is RaggedAxis:
[348](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/batching.py:348) new_aval = core.mapped_aval(
[349](https://vscode-remote+dev-002dcontainer-002b7b22686f737450617468223a225c5c5c5c77736c2e6c6f63616c686f73745c5c5562756e74755c5c686f6d655c5c67616e67656c6f7474695c5c70726f6a656374735c5c74656e736f72666c6178222c226c6f63616c446f636b6572223a66616c73652c22636f6e66696746696c65223a7b22246d6964223a312c2270617468223a222f686f6d652f67616e67656c6f7474692f70726f6a656374732f74656e736f72666c61782f2e646576636f6e7461696e65722e6a736f6e222c22736368656d65223a227673636f64652d66696c65486f7374227d7d.vscode-resource.vscode-cdn.net/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/batching.py:349) aval.shape[self.batch_dim.stacked_axis], self.batch_dim.stacked_axis, aval)
IndexError: tuple index out of range
### Steps to reproduce:
```python
from flax import nnx
import jax
from functools import partial
class MLP(nnx.Module):
def __init__(self, din: int, dmid: int, dout: int, *, rngs: nnx.Rngs):
self.linear1 = nnx.Linear(din, dmid, rngs=rngs)
self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
self.bn = nnx.BatchNorm(dmid, rngs=rngs)
self.linear2 = nnx.Linear(dmid, dout, rngs=rngs)
def __call__(self, x: jax.Array):
x = nnx.gelu(self.dropout(self.bn(self.linear1(x))))
return self.linear2(x)
@partial(nnx.vmap, axis_size=5)
def create_model(rngs: nnx.Rngs):
return MLP(10, 32, 10, rngs=rngs)
model = create_model(nnx.Rngs(0))
```
Contributor guide
Assessment
This issue has not been assessed yet.