ToLinen is not hashable (Linen modules are)
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Linen modules are washable, so I would expect `nnx.bridge.ToLinen` to be as well.
```python
In [1]: from flax import linen as nn, nnx
In [2]: import jax
In [3]: model = nnx.bridge.ToLinen(nnx.Linear, args=(32, 64))
In [4]: hash(model)
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
File ~/Documents/pythonenvs/netket/python-3.11.2/lib/python3.11/site-packages/flax/linen/module.py:726, in _wrap_hash..wrapped(self)
725 try:
--> 726 hash_value = hash_fn(self)
727 except TypeError as exc:
File :3, in __hash__(self)
TypeError: unhashable type: 'dict'
The above exception was the direct cause of the following exception:
TypeError Traceback (most recent call last)
Cell In[4], line 1
----> 1 hash(model)
File ~/Documents/pythonenvs/netket/python-3.11.2/lib/python3.11/site-packages/flax/linen/module.py:728, in _wrap_hash..wrapped(self)
726 hash_value = hash_fn(self)
727 except TypeError as exc:
--> 728 raise TypeError(
729 'Failed to hash Flax Module. '
730 'The module probably contains unhashable attributes. '
731 f'Module={self}'
732 ) from exc
733 return hash_value
TypeError: Failed to hash Flax Module. The module probably contains unhashable attributes. Module=ToLinen(
# attributes
nnx_class = Linear
args = (32, 64)
kwargs = {}
skip_rng = False
)
```
This is problematic because I cannot use `ToLinen` in lieu of standard linen modules which I pass as static arguments jax jax.jit functions.
Contributor guide
Assessment
This issue has not been assessed yet.