google / google/flax

Cannot assign arrays to dataclass fields in `nnx`

Open
#3,627 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

When instantiating a `nnx.dataclass` module, if the input to a `param_field` (or any `variable_field` actually) is a jax Array, a ValueError is raised because the value is assigned to the module without being wrapped into the `nnx.Param` class.

```python
import jax
import jax.numpy as jnp
from flax.experimental import nnx

@nnx.dataclass
class Foo(nnx.Module):
x: jax.Array = nnx.param_field()

foo = Foo(jnp.array(0.2))
```

I would expect that the input is wrapped into `nnx.Param` before being assigned to the module. Same as is works for e.g. integers or floats.

### Logs, error messages, etc:

```
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
/users/fzanetta/pyprojects/GPJax/_debug/variables.ipynb Cell 2 line 5
[1] @nnx.dataclass
[2] class Foo(nnx.Module):
[3] x: jax.Array = nnx.param_field()
Foo(jnp.array(0.2))

File ~/.cache/pypoetry/virtualenvs/gpjax-giodCE1Q-py3.10/lib/python3.10/site-packages/flax/experimental/nnx/nnx/module.py:150, in ModuleMeta.__call__(self, *args, **kwargs)
149 def __call__(self, *args: Any, **kwargs: Any) -> Any:
--> 150 return self._meta_call(*args, **kwargs)

File ~/.cache/pypoetry/virtualenvs/gpjax-giodCE1Q-py3.10/lib/python3.10/site-packages/flax/experimental/nnx/nnx/module.py:155, in ModuleMeta._meta_call(cls, *args, **kwargs)
153 module = cls.__new__(cls, *args, **kwargs)
154 vars(module)['_module__state'] = ModuleState()
--> 155 module.__init__(*args, **kwargs)
157 if dataclasses.is_dataclass(module):
158 if isinstance(module, _HasSetup):

File :3, in __init__(self, x)

File ~/.cache/pypoetry/virtualenvs/gpjax-giodCE1Q-py3.10/lib/python3.10/site-packages/flax/experimental/nnx/nnx/module.py:207, in Module.__setattr__(self, name, value)
206 def __setattr__(self, name: str, value: Any) -> None:
--> 207 self._setattr(name, value)

File ~/.cache/pypoetry/virtualenvs/gpjax-giodCE1Q-py3.10/lib/python3.10/site-packages/flax/experimental/nnx/nnx/module.py:232, in Module._setattr(self, name, value)
230 else:
231 if isinstance(value, (jax.Array, np.ndarray, State)):
--> 232 raise ValueError(
233 f\"Trying to assign a '{type(value).__name__}' to the Module\"
234 f\" attribute '{name}'. This is not supported. Non-hashable \"
235 'objects are not valid static state in JAX. Please wrap '
236 'the value in a Variable type instead.'
237 )
238 vars_dict[name] = value

ValueError: Trying to assign a 'ArrayImpl' to the Module attribute 'x'. This is not supported. Non-hashable objects are not valid static state in JAX. Please wrap the value in a Variable type instead."
}
```

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.