0.12.7

What's Changed

• fix transform_metadata by @copybara-service[bot] in #5368
• Add 'opt_sharding' option for Variables by @samanklesaria in #5231
• Move orbax-export to optional deps by @samanklesaria in #5353
• Add nnx.make_jaxpr. by @copybara-service[bot] in #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 #5373
• Add with_rngs function by @samanklesaria in #5383
• Add iter_module_children() to iterate over module children only. by @copybara-service[bot] in #5302
• Return updates PyTree from nnx.Optimizer.update by @anshuljoshi in #5214
• Add deprecations file by @samanklesaria in #5382
• Fix doctests for optimizer update by @samanklesaria in #5386
• Deprecate the keep_rngs argument in recurrent.py by @copybara-service[bot] in #5385
• add nnx.prefix by @copybara-service[bot] in #5387
• add Rngs.broadcast by @copybara-service[bot] in #5390
• expose nnx pooling layers in API docs by @copybara-service[bot] in #5392
• improve with_rngs by @copybara-service[bot] in #5391
• Removed bonsai from docs by @vfdev-5 in #5380
• [pmap] In-line definitions of jax.device_put_sharded and jax.device_put_replicated. by @copybara-service[bot] in #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 #5397
• Rename pure -> as_pure by @samanklesaria in #5389
• change defaults for Dropout and BatchNorm by @copybara-service[bot] in #5299
• Fix CI by @cgarciae in #5400
• Avoid deprecated pattern of constructing a dtype from an array. by @copybara-service[bot] in #5402
• Fix for change to jax.Array.__contains__ by @copybara-service[bot] in #5404
• forward fix simple transforms by default by @copybara-service[bot] in #5393
• improve nnx.compat by @copybara-service[bot] in #5406
• Fixed broken url in the sow warning by @vfdev-5 in #5377
• Updates in gemma example by @vfdev-5 in #4913
• [JAX] Fix code that compares PartitionSpecs with tuples. by @copybara-service[bot] in #5410
• Add flag for variables in nnx.merge by @samanklesaria in #5381
• fix Rngs.getattr returning default stream for dunder names (#5375) by @mohsinm-dev in #5379
• Add Optimization Cookbook by @samanklesaria in #5117
• Update pytree guide to use nnx.capture by @samanklesaria in #5370
• Add diffusion example ported from Jax AI stack by @samanklesaria in #5403
• Rename abstract_with_sharding to as_abstract by @samanklesaria in #5388
• Avoid deprecated JAX API by @copybara-service[bot] in #5414
• improve check_prefix by @copybara-service[bot] in #5407
• use State instead of GraphState by @copybara-service[bot] in #5417
• add nnx.unpack by @copybara-service[bot] in #5418
• Added kwargs to call to pass kwargs to basemodule in nnx.LoRA by @vfdev-5 in #5421
• Avoid usage of deprecated jax.core APIs. by @copybara-service[bot] in #5430
• update version to 0.12.7 by @copybara-service[bot] in #5429
• Deprecate Variable hooks in favor of custom Variable types. by @copybara-service[bot] in #5415
• Fix doubled 'the the' typos across multiple files by @SAY-5 in #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 #5435

New Contributors

• @anshuljoshi made their first contribution in #5214
• @SAY-5 made their first contribution in #5422

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