Flax Files

What's Changed

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

New Contributors

• @anshuljoshi made their first contribution in https://github.com/google/flax/pull/5214
• @SAY-5 made their first contribution in https://github.com/google/flax/pull/5422

Full Changelog: https://github.com/google/flax/compare/v0.12.6...v0.12.7