0.12.7

[cgarciae]

What's Changed

• fix transform_metadata by @copybara-service[bot] in fix transform_metadata #5368
• Add 'opt_sharding' option for Variables by @samanklesaria in Add 'opt_sharding' option for Variables #5231
• Move orbax-export to optional deps by @samanklesaria in Move orbax-export to optional deps #5353
• Add nnx.make_jaxpr. by @copybara-service[bot] in Add nnx.make_jaxpr. #5378
• Replaced custom PriorityStr with standard tuple sorting in _graph_node_flatten. Python 3.13.11 optimized str comparison bypasses custom subclasses (causing TypeError), and used getattr in _setattr because unflattening order becomes non-deterministic (preventing AttributeError). by @copybara-service[bot] in Replaced custom PriorityStr with standard tuple sorting in _graph_node_flatten. Python 3.13.11 optimized str comparison bypasses custom subclasses (causing TypeError), and used getattr in _setattr because unflattening order becomes non-deterministic (preventing AttributeError). #5373
• Add with_rngs function by @samanklesaria in Add with_rngs function #5383
• Add iter_module_children() to iterate over module children only. by @copybara-service[bot] in Add iter_module_children() to iterate over module children only. #5302
• Return updates PyTree from nnx.Optimizer.update by @anshuljoshi in Return updates PyTree from nnx.Optimizer.update #5214
• Add deprecations file by @samanklesaria in Add deprecations file #5382
• Fix doctests for optimizer update by @samanklesaria in Fix doctests for optimizer update #5386
• Deprecate the keep_rngs argument in recurrent.py by @copybara-service[bot] in Deprecate the keep_rngs argument in recurrent.py #5385
• add nnx.prefix by @copybara-service[bot] in add nnx.prefix #5387
• add Rngs.broadcast by @copybara-service[bot] in add Rngs.broadcast #5390
• expose nnx pooling layers in API docs by @copybara-service[bot] in expose nnx pooling layers in API docs #5392
• improve with_rngs by @copybara-service[bot] in improve with_rngs #5391
• Removed bonsai from docs by @vfdev-5 in Removed bonsai from docs #5380
• [pmap] In-line definitions of jax.device_put_sharded and jax.device_put_replicated. by @copybara-service[bot] in [pmap] In-line definitions of jax.device_put_sharded and jax.device_put_replicated. #5374
• Fix the wrong names for the hooks in nnx.with_metadata, fixed such that names match the proper hook names in Variable class. by @copybara-service[bot] in Fix the wrong names for the hooks in nnx.with_metadata, fixed such that names match the proper hook names in Variable class. #5397
• Rename pure -> as_pure by @samanklesaria in Rename pure -> as_pure #5389
• change defaults for Dropout and BatchNorm by @copybara-service[bot] in change defaults for Dropout and BatchNorm #5299
• Fix CI by @cgarciae in Fix CI #5400
• Avoid deprecated pattern of constructing a dtype from an array. by @copybara-service[bot] in Avoid deprecated pattern of constructing a dtype from an array. #5402
• Fix for change to jax.Array.__contains__ by @copybara-service[bot] in Fix for change to jax.Array.__contains__ #5404
• forward fix simple transforms by default by @copybara-service[bot] in forward fix simple transforms by default #5393
• improve nnx.compat by @copybara-service[bot] in improve nnx.compat #5406
• Fixed broken url in the sow warning by @vfdev-5 in Fixed broken url in the sow warning #5377
• Updates in gemma example by @vfdev-5 in Updates in gemma example #4913
• [JAX] Fix code that compares PartitionSpecs with tuples. by @copybara-service[bot] in [JAX] Fix code that compares PartitionSpecs with tuples. #5410
• Add flag for variables in nnx.merge by @samanklesaria in Add flag for variables in nnx.merge #5381
• fix Rngs.getattr returning default stream for dunder names (nnx.Rngs.__*__ returns the default stream, which causes confusing breakage in combination with other packages #5375) by @mohsinm-dev in fix Rngs.__getattr__ returning default stream for dunder names (#5375) #5379
• Add Optimization Cookbook by @samanklesaria in Add Optimization Cookbook #5117
• Update pytree guide to use nnx.capture by @samanklesaria in Update pytree guide to use nnx.capture #5370
• Add diffusion example ported from Jax AI stack by @samanklesaria in Add diffusion example ported from Jax AI stack #5403
• Rename abstract_with_sharding to as_abstract by @samanklesaria in Rename abstract_with_sharding to as_abstract #5388
• Avoid deprecated JAX API by @copybara-service[bot] in Avoid deprecated JAX API #5414
• improve check_prefix by @copybara-service[bot] in improve check_prefix #5407
• use State instead of GraphState by @copybara-service[bot] in use State instead of GraphState #5417
• add nnx.unpack by @copybara-service[bot] in add nnx.unpack #5418
• Added kwargs to call to pass kwargs to basemodule in nnx.LoRA by @vfdev-5 in Added kwargs to call to pass kwargs to basemodule in nnx.LoRA #5421
• Avoid usage of deprecated jax.core APIs. by @copybara-service[bot] in Avoid usage of deprecated jax.core APIs. #5430
• update version to 0.12.7 by @copybara-service[bot] in update version to 0.12.7 #5429
• Deprecate Variable hooks in favor of custom Variable types. by @copybara-service[bot] in Deprecate Variable hooks in favor of custom Variable types. #5415
• Fix doubled 'the the' typos across multiple files by @SAY-5 in Fix doubled 'the the' typos across multiple files #5422
• Disallow P.UNCONSTRAINED in out_shardings in jit. This leads to simplifications in how we lower shardings to stableHLO. by @copybara-service[bot] in Disallow P.UNCONSTRAINED in out_shardings in jit. This leads to simplifications in how we lower shardings to stableHLO. #5435

New Contributors

• @anshuljoshi made their first contribution in Return updates PyTree from nnx.Optimizer.update #5214
• @SAY-5 made their first contribution in Fix doubled 'the the' typos across multiple files #5422

Full Changelog: v0.12.6...v0.12.7

---

This discussion was created from the release 0.12.7.