google / google/flax

ToLinen is not hashable (Linen modules are)

Open
#4,156 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

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

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.