v0.12.1

[cgarciae]

Deprecations

Variable.value

Variable.value is now deprecated. Consider the following example:

import jax.numpy as jnp
import jax
from flax import nnx

my_param = nnx.Param({'a': 0.0})

@nnx.jit
def f(m):
    m.value['a'] = 1.0
    return m

Running f(my_param) produces Param(value={'a': 0.0}), not Param(value={'a': 1.0}) as before. This is because getting the value parameter new returns a copy of the pytree values (like dict / list). Instead, use the __setitem__ method to update the value:

@nnx.jit
def f(m):
    m['a'] = 1.0
    return m

nnx.Data and nnx.Static

nnx.Data and nnx.Static annotations are now deprecated. To create nnx.Pytree or nnx.Module dataclasses use the new nnx.dataclass with nnx.data and nnx.static as field descriptors.

# old
@dataclasses.dataclass
class Foo(nnx.Pytree):
  a: nnx.Data[int]
  b: nnx.Static[str]

# new
@nnx.dataclass
class Foo(nnx.Pytree):
  a: int = nnx.data()
  b: str = nnx.static()

Pull Requests

• Clarify *Norm layer docstrings: axis_index_groups is unused under SPMD jit. by @copybara-service[bot] in Clarify *Norm layer docstrings: axis_index_groups is unused under SPMD jit. #4940
• Move ArrayRef creation to the end of Variable creation by @IvyZX in Move ArrayRef creation to the end of Variable creation #4980
• clean up jax.Ref-related names by @copybara-service[bot] in clean up jax.Ref-related names #4988
• Add compute_flops and compute_vjp_flops options to nnx.tabulate by @samanklesaria in Add compute_flops and compute_vjp_flops options to nnx.tabulate #4948
• Fix nnx.tabulate crash with empty dict/None values (fixes Edge case in nnx.tabulate when nnx.Module stores empty dictionary #4889) by @mohsinm-dev in Fix nnx.tabulate crash with empty dict/None values (fixes #4889) #4891
• Future-proof imports of jax.new_ref / jax.Ref. by @copybara-service[bot] in Future-proof imports of jax.new_ref / jax.Ref. #4986
• Use jnp.stack instead of np.stack in flax.training.common_utils.stack_forest by @vfdev-5 in Use jnp.stack instead of np.stack in flax.training.common_utils.stack_forest #4991
• Fixed broken nnx.statelib.diff by @vfdev-5 in Fixed broken nnx.statelib.diff #4992
• Implemented spectral norm in NNX by @mattbahr in Implemented spectral norm in NNX #4623
• Improve Variable.{get,set}_metadata by @cgarciae in Improve Variable.{get,set}_metadata #4985
• Move iter_children and iter_modules to functions by @samanklesaria in Move iter_children and iter_modules to functions #4961
• Avoid install, import, or tests with tensorflow-text under Python 3.13+. by @jburnim in Avoid install, import, or tests with tensorflow-text under Python 3.13+. #5001
• disallow setting metadata through settattr by @cgarciae in disallow setting metadata through settattr #4993
• Use sphinx 6.2+ for docs, which works with Python 3.13. by @jburnim in Use sphinx 6.2+ for docs, which works with Python 3.13. #5009
• Removed kernel_init/bias_init atttributes from popular layers by @vfdev-5 in Removed kernel_init/bias_init atttributes from popular layers #4998
• Migrate from jax.experimental.enable_x64 to jax.enable_x64. by @copybara-service[bot] in Migrate from jax.experimental.enable_x64 to jax.enable_x64. #5011
• Add Rngs KeylessInitializers by @cgarciae in Add Rngs KeylessInitializers #5017
• optimize scan transpositions by @cgarciae in optimize scan transpositions #5015
• Variable refactor by @cgarciae in Variable refactor #5006
• Remove invalid gymnasium dependency in pyproject.toml by @IvyZX in Remove invalid gymnasium dependency in pyproject.toml #5016
• Use jax.shard_map in flax by @copybara-service[bot] in Use jax.shard_map in flax #5020
• use jax.shard_map by @copybara-service[bot] in use jax.shard_map #5018
• Fix formatting in PR template checklist by @rapsealk in Fix formatting in PR template checklist #5024
• Fixed attribute visualization in treescope_repr by @vfdev-5 in Fixed attribute visualization in __treescope_repr__ #5022
• feat: add nnx.set_metadata to in-place change metadata of the state variables of nnx.Modules by @pfackeldey in feat: add nnx.set_metadata to in-place change metadata of the state variables of nnx.Modules #5007
• Update README to use fully qualified nnx.Linear in example by @rapsealk in Update README to use fully qualified nnx.Linear in example #5023
• Fix nnx tabulate variable hooks by @mohsinm-dev in Fix nnx tabulate variable hooks #5008
• python 3.13 support by @cgarciae in python 3.13 support #4987
• Added a note in nnx.jit about arg donation by @vfdev-5 in Added a note in nnx.jit about arg donation #5031
• Add flip doc link to eager sharding error message by @IvyZX in Add flip doc link to eager sharding error message #5033
• fix reseed for abstract values by @cgarciae in fix reseed for abstract values #5034
• Deduplicate Variable nodes in iter_graph and eliminate recursion. by @copybara-service[bot] in Deduplicate Variable nodes in iter_graph and eliminate recursion. #5035
• Support for python 3.14 by @vfdev-5 in Support for python 3.14 #5032
• [docs] Exposed more helper functions/classes in state.rst by @vfdev-5 in [docs] Exposed more helper functions/classes in state.rst #5037
• Copybara import of the project: by @copybara-service[bot] in Copybara import of the project: #5041
• Internal change by @copybara-service[bot] in Internal change #5048
• filter grad state in nnx.Optimizer by @copybara-service[bot] in filter grad state in nnx.Optimizer #5049
• Add NNX WeightNorm (update of [nnx] Add NNX WeightNorm. #4568) by @samanklesaria in Add NNX WeightNorm (update of #4568) #5043
• Fix shard_map documentation link in compilation.py by @vfdev-5 in Fix shard_map documentation link in compilation.py #5038
• Fix ValueError when nnx.jit is used with nnx.custom_vjp by @samanklesaria in Fix ValueError when nnx.jit is used with nnx.custom_vjp #5045
• Recursive map by @chapman20j in Recursive map #5042
• Convert linen pytorch guide to nnx by @samanklesaria in Convert linen pytorch guide to nnx #4999
• Set Mode with Tests by @chapman20j in Set Mode with Tests #5056
• Fixing Optimizer docstring - fixing Documentation Inconsistency regarding flax.nnx.optimizer.Optimizer #5060 by @Lucas-Fernandes-Martins in Fixing Optimizer docstring - fixing #5060 #5061
• Update tutorial examples to thread explicit RNGs by @samanklesaria in Update tutorial examples to thread explicit RNGs #4975
• Fix NNX jit static args with in_shardings issue NNX JIT fails with static args and in shardings #4989 by @mohsinm-dev in Fix NNX jit static args with in_shardings issue #4989 #4996
• support explicit sharding in eager sharding by @cgarciae in support explicit sharding in eager sharding #5070
• Added missing LayerNorm test case into TestLayersSameGraph by @vfdev-5 in Added missing LayerNorm test case into TestLayersSameGraph #5076
• fix main by @cgarciae in fix main #5081
• docs: Document allow_duplicates argument of nnx.to_arrays. by @dan-zheng in docs: Document allow_duplicates argument of nnx.to_arrays. #5083
• add promote_dtype to all standard layers by @cgarciae in add promote_dtype to all standard layers #5080
• add nnx.dataclass by @cgarciae in add nnx.dataclass #5066
• Expand ConvTranspose padding documentation by @samanklesaria in Expand ConvTranspose padding documentation #4990
• Added kernel_metadata/bias_metadata args to nnx layers by @vfdev-5 in Added kernel_metadata/bias_metadata args to nnx layers #5074
• Add nnx.use_eager_sharding context manager by @samanklesaria in Add nnx.use_eager_sharding context manager #5079
• fix main by @cgarciae in fix main #5090
• Adding set_mode_info by @chapman20j in Adding set_mode_info #5071
• Fixed nnx.scan with carry as pytree and sow by @vfdev-5 in Fixed nnx.scan with carry as pytree and sow #5073
• Fix bound method auto-unbinding for NNX transforms by @mohsinm-dev in Fix bound method auto-unbinding for NNX transforms #5055
• deprecate Variable.value by @cgarciae in deprecate Variable.value #5052
• Add eq for variables by @samanklesaria in Add __eq__ for variables #5084
• Fixed deprecated .value usage failing CI tests by @vfdev-5 in Fixed deprecated .value usage failing CI tests #5097
• update jax minver to 0.8.1 by @cgarciae in update jax minver to 0.8.1 #5095

New Contributors

• @samanklesaria made their first contribution in Add compute_flops and compute_vjp_flops options to nnx.tabulate #4948
• @jburnim made their first contribution in Avoid install, import, or tests with tensorflow-text under Python 3.13+. #5001
• @rapsealk made their first contribution in Fix formatting in PR template checklist #5024
• @pfackeldey made their first contribution in feat: add nnx.set_metadata to in-place change metadata of the state variables of nnx.Modules #5007
• @chapman20j made their first contribution in Recursive map #5042
• @Lucas-Fernandes-Martins made their first contribution in Fixing Optimizer docstring - fixing #5060 #5061

Full Changelog: v0.12.0...v0.12.1

---

This discussion was created from the release v0.12.1.