Skip to content

fix(nnx): make ToLinen hashable for JAX JIT (#4156) - #5562

Open
Ultron09 wants to merge 1 commit into
google:mainfrom
Ultron09:fix/tolinen-hashable
Open

fix(nnx): make ToLinen hashable for JAX JIT (#4156)#5562
Ultron09 wants to merge 1 commit into
google:mainfrom
Ultron09:fix/tolinen-hashable

Conversation

@Ultron09

Copy link
Copy Markdown

What does this PR do?

Fixes #4156

Standard Linen modules are hashable dataclasses, allowing them to be passed as static arguments to jax.jit functions or used in caches. However, flax.nnx.bridge.ToLinen could fail with TypeError: unhashable type: 'dict' or unhashable sequence errors if instantiated with mutable args (e.g. list) or kwargs (e.g. standard dict).

This PR:

  1. Adds __post_init__ in ToLinen to ensure self.args is converted to a tuple and self.kwargs is converted to a flax.core.FrozenDict.
  2. Adds a comprehensive regression test test_tolinen_hashable in tests/nnx/bridge/wrappers_test.py covering direct ToLinen, to_linen, and dynamic to_linen_class modules.

Checklist

  • This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).
  • This change is discussed in a Github issue/discussion.
  • The documentation and docstrings adhere to the documentation guidelines.
  • This change includes necessary high-coverage tests. (No quality testing = no merge!)

cc @cgarciae @vfdev-5 for review.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ToLinen is not hashable (Linen modules are)

2 participants