Releases

@copybara-service

Reverted changes to traverse utils.

What's Changed

rollback of breaking change by @copybara-service[bot] in #5629
Release Flax 0.12.12 by @copybara-service[bot] in #5630

Full Changelog: v0.12.11...v0.12.12

@samanklesaria

What's Changed

Fix mypy hijax error by @samanklesaria in #5580
Fix nnx.cond tracing cache misses by caching SimpleCondFn wrappers by @chenkuanliao in #5518
Add nnx_sow_capture_noop config flag and capture_noop option to Module.sow. by @copybara-service[bot] in #5607
Replace # pytype: disable suppressions with # pyrefly: ignore by @copybara-service[bot] in #5612
Fix Flax GitHub Actions CI (pytest DeprecationWarning and pytype in examples/gemma/train.py). by @copybara-service[bot] in #5616
Don't pass linen_meta_type argument when creating AxisMetadata subclasses. by @copybara-service[bot] in #4933
Add capture_noop option to Module.perturb. by @copybara-service[bot] in #5621
Add sepreformer example by @samanklesaria in #5472
Add multimodal llm example by @samanklesaria in #5488
Store nnx.WeightNorm scales as Param so they are trainable by @IvyZX in #5623
[nnx] Make nnx.Sequential iterable and filter **kwargs per layer. by @copybara-service[bot] in #5625
Release Flax 0.12.11 by @copybara-service[bot] in #5627
fix(nnx): skip Sequence modules in capture to prevent mixed key types (fixes #5571) by @jdymitarai in #5582
Fix NNX MultiHeadAttention is_causal mask during decode by @Anu27n in #5592
fix(traverse_util): raise ValueError when flatten_dict separator conflicts with keys by @mayuriphad in #5593

New Contributors

@jdymitarai made their first contribution in #5582
@Anu27n made their first contribution in #5592
@mayuriphad made their first contribution in #5593

Full Changelog: v0.12.10...v0.12.11

@copybara-service

What's Changed

verify Flax RoPE using mathematical definition instead of PyTorch by @copybara-service[bot] in #5558
Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive by @copybara-service[bot] in #5566
[No functional change] Delete VJPHiPrimitive and replace with HiPrim. This is a simple replace. by @copybara-service[bot] in #5568
suppress pyrefly errors related to JAX scalar types by @copybara-service[bot] in #5586
Release Flax 0.12.10 and use attribute lookup for HiPrim fallback by @copybara-service[bot] in #5597

Full Changelog: v0.12.9...v0.12.10

@chenkuanliao

What's Changed

Add vmap reset() regression test for nnx.metrics (#5483) by @chenkuanliao in #5491
Validate dropout_rate in dot_product_attention by @mohsinm-dev in #5503
Add RoPE embeddings by @samanklesaria in #5481
implement axes_scan's broadcast semantics via dce by @mattjj in #5540
Refactor NNX Variable to remove legacy JAX QDD dependencies. by @copybara-service[bot] in #5541
[typing] add assertion to avoid pyrefly error by @copybara-service[bot] in #5546
add pyrefly suppressions for jax.jit annotation change by @copybara-service[bot] in #5547
Early exit in Pytree._setattr for data descriptors. by @copybara-service[bot] in #5549
replace lazy_init implementation to avoid jax partial_eval by @mattjj in #5551
Make Flax thread-safe for multi-threaded tracing. by @copybara-service[bot] in #5543
release Flax 0.12.9 by @copybara-service[bot] in #5557

New Contributors

@chenkuanliao made their first contribution in #5491

Full Changelog: v0.12.8...v0.12.9

@vfdev-5

What's Changed

Fix intermediates for Gemma example by @vfdev-5 in #5409
Update view guide to mention with_ and as_ prefixes by @samanklesaria in #5416
Reorganize examples by @samanklesaria in #5424
[Docs] migration guide "From PyTorch to JAX and Flax" by @vfdev-5 in #5408
Rename vars_as to with_vars by @samanklesaria in #5384
Add machine translation example by @samanklesaria in #5431
Updated gemma docs by @vfdev-5 in #5427
Reenabled examples tests by @vfdev-5 in #5098
fix Pytree typing issue by @copybara-service[bot] in #5449
improve Optimizer by @copybara-service[bot] in #5450
Added tests for flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate by @vfdev-5 in #5101
Fix docstrings in graphlib under graph=False by @samanklesaria in #5443
Port Jax AI stack minigpt example by @samanklesaria in #5405
Add fix for msgpack serialization of nnx.List structures by @samanklesaria in #5446
Added out_sharding and qkv_sharding to MHA call method by @vfdev-5 in #5434
add nnx.EMA by @copybara-service[bot] in #5451
Update mnist tutorial to use tensorboard by @samanklesaria in #5425
Fix rnglib doctests for graph_updates=False default by @samanklesaria in #5442
Fix update doctest for graph_updates=False default by @samanklesaria in #5440
Fix doctests in transforms for graph_updates=False default by @samanklesaria in #5441
Update the RNG guide with split and without StateAxes by @samanklesaria in #5363
Add data loaders tutorial by @samanklesaria in #5454
add storage backend type in saving telemetry. by @copybara-service[bot] in #5465
optimize transform updates by @copybara-service[bot] in #5453
Added image segmentation example by @vfdev-5 in #5463
Fixed gemma example sampler test test_models_from_kaggle by @vfdev-5 in #5477
fix MultiMetric GraphDef not hashable under jit by @mohsinm-dev in #5492
Added ViT training example by @vfdev-5 in #5455
Fixed CI issues: doctest, test examples by @vfdev-5 in #5506
Add Dependabot configuration for dependency updates by @vfdev-5 in #5507
optimize jit_partial by @copybara-service[bot] in #5476
update to v0.12.8 by @copybara-service[bot] in #5509
Update tensorflow dependency versions in pyproject.toml by @vfdev-5 in #5511
Fixed typos in docs_nnx/migrating/pytorch_to_jax_flax.rst by @vfdev-5 in #5479
fix nnx promote_dtype docstring by @mohsinm-dev in #5504
replace bare asserts with ValueError in NNX attention validation by @mohsinm-dev in #5505
Added Object Detection example with minimal DETR by @vfdev-5 in #5480
Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES by @copybara-service[bot] in #5516
nnx.clone: add arrays parameter to also copy underlying buffers by @dparikh79 in #5475
Automated Code Change by @copybara-service[bot] in #5520
Use less precise numbers in linen doctest by @samanklesaria in #5526

New Contributors

@dparikh79 made their first contribution in #5475

Full Changelog: v0.12.7...v0.12.8

@copybara-service

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

@copybara-service

What's Changed

Annotate wrapped_fn's self argument in decorator_lift_transform for trace-ability by @copybara-service[bot] in #5298
add with_attributes by @copybara-service[bot] in #5308
Remove optax pin by @samanklesaria in #5292
Removed redundant code checking wrt in nnx.Optimizer by @vfdev-5 in #5226
Fix PyTreeNode + Generic losing parameters when Generic is last in bases by @mohsinm-dev in #5237
add recursive_map example by @copybara-service[bot] in #5311
Add sharding propagation support in nnx.eval_shape (clone of #5111) by @samanklesaria in #5247
fix eval_shape's _to_variable by @copybara-service[bot] in #5316
support out_sharding mapping in set_metadata by @copybara-service[bot] in #5312
feat(nnx): add GQA support to MultiHeadAttention by @ayulockedin in #5259
Allow DenyList to be compared. by @copybara-service[bot] in #5322
add graph_updates argument to jit by @copybara-service[bot] in #5317
add graph_updates argument in shard_map by @copybara-service[bot] in #5319
Remove jax/tools/colab_tpu.py. by @copybara-service[bot] in #5324
add graph_updates to vmap by @copybara-service[bot] in #5320
Add split method to RngStream by @samanklesaria in #5270
add graph_updates for scan by @copybara-service[bot] in #5327
add graph_updates to while_loop by @copybara-service[bot] in #5328
Added support for data masking in Average, Accuracy and MultiMetric by @vfdev-5 in #5326
add graph_updates to fori_loop by @copybara-service[bot] in #5329
add graph_updates to pmap, grad, and value_and_grad by @copybara-service[bot] in #5330
add tree-mode-nnx FLIP by @copybara-service[bot] in #5310
Copybara import of the project: by @copybara-service[bot] in #5331
add graph_updates to remat by @copybara-service[bot] in #5336
add graph_updates to eval_shape and checkify by @copybara-service[bot] in #5338
add compat module by @copybara-service[bot] in #5340
add more tests that check for consistent aliasing in transforms by @copybara-service[bot] in #5341
don't allow Variable mutation in custom_vjp on differentiable arguments by @copybara-service[bot] in #5342
support hijax and ref Variables in simple transforms by @copybara-service[bot] in #5345
allow in/out_axes when graph_updates=False by @copybara-service[bot] in #5323
add transform_metadata transform by @copybara-service[bot] in #5346
clean up custom_vjp's graph_updates=False section by @copybara-service[bot] in #5350
improve error messages for tree mode duplicates check by @copybara-service[bot] in #5347
Introduce manual_type: ManualAxisType parameter on ShapedArray to track varying/unreduced/reduced and remove vma: frozenset parameter. by @copybara-service[bot] in #5339
remove aliases from nnx.graph in favor of nnx.compat by @copybara-service[bot] in #5348
check aliases on all transform args and simplify apply_variable_updates by @copybara-service[bot] in #5349
Add intermediate value captures (extends #4925) by @samanklesaria in #5257
fix jit_partial lower and compile by @copybara-service[bot] in #5355
Added support for data masking in Average, Accuracy and MultiMetric by @vfdev-5 in #5332
Do a few more cleanups after pmap_shmap merge by @copybara-service[bot] in #5358
update to version 0.12.6 by @copybara-service[bot] in #5356
Do a few more cleanups after pmap_shmap merge by @copybara-service[bot] in #5359
Do a few more cleanups after the pmap shmap merge deletion. by @copybara-service[bot] in #5361
simplify SimpleScan by @copybara-service[bot] in #5357
improve nnx.Dict error handling by @copybara-service[bot] in #5362
add graph node in prefix checks by @copybara-service[bot] in #5365
improve aliasing error msg by @copybara-service[bot] in #5364
Generalize out_sharding to work with NamedSharding and Format by @samanklesaria in #5246
add nnx.map and nnx.abstract_with_sharding by @copybara-service[bot] in #5366

Full Changelog: v0.12.5...v0.12.6

@copybara-service

What's Changed

maintain data/static definition in split / state by @copybara-service[bot] in #5243
Add graph=False support for nnx.grad and nnx.value_and_grad by @copybara-service[bot] in #5240
Add tree-mode support for nnx.remat by @copybara-service[bot] in #5242
Add tree-mode support for nnx.vmap by @copybara-service[bot] in #5250
Add tree-mode arg error handling by @copybara-service[bot] in #5251
[flax:benchmarks] Fix flax benchmarks. by @copybara-service[bot] in #5241
Improve unregistered data detection by @copybara-service[bot] in #5249
Add graph argument to split, state, graphdef, clone by @copybara-service[bot] in #5258
Add nnx_graph_mode config flag by @copybara-service[bot] in #5261
Add tree-mode support to nnx.{cond,switch,eval_shape,checkify} by @copybara-service[bot] in #5252
Add tree-mode support to nnx.scan by @copybara-service[bot] in #5262
Add tree-mode support for nnx.custom_vjp by @copybara-service[bot] in #5264
Add tree-mode support to nnx.{while,fori}_loop by @copybara-service[bot] in #5263
add tree-mode support for pmap by @copybara-service[bot] in #5265
Add tree-mode support to iter_graph by @copybara-service[bot] in #5266
Add tree-mode support for nnx.recursive_map, nnx.view, and nnx.view_info by @copybara-service[bot] in #5267
Use linen_vars_to_nnx_attrs in ToLinen variable restoration by @copybara-service[bot] in #5272
Add tree-mode-only nnx.{vjp,jvp,jit_partial} by @copybara-service[bot] in #5268
Use linen_vars_to_nnx_attrs in ToLinen variable restoration by @copybara-service[bot] in #5273
fix nn.vmap's _broadcast_prefix_tree by @copybara-service[bot] in #5282
fix nn.vmap's split_fn by @copybara-service[bot] in #5286
[pmap] Remove jax.config.pmap_shmap_merge. by @copybara-service[bot] in #5289
add set_graph_mode by @copybara-service[bot] in #5269
add error handling to iter_children by @copybara-service[bot] in #5296
expose and update view guide by @cgarciae in #5294
Remove chex dependency by @samanklesaria in #5295
improve error messages for tree mode errors by @copybara-service[bot] in #5301
remove experimental Variable.mutable property by @copybara-service[bot] in #5304
update version to 0.12.5 by @copybara-service[bot] in #5305
simplify nnx.display for Variable by @copybara-service[bot] in #5306
simplify set_view API by @copybara-service[bot] in #5303

Full Changelog: v0.12.4...v0.12.5

@cgarciae

This release fixes an issue with nnx.List and nnx.Sequential passing all its elements as static when being treated as a pytree.

What's Changed

fix List and improve Pytree by @cgarciae in #5072
Copybara import of the project: by @copybara-service[bot] in #5211
Fixed mask sharding if inputs is sharded by @vfdev-5 in #5212
nnx.view by @chapman20j in #5204
bonsai link by @chapman20j in #5220
Add out_sharding argument to Embed layer call by @samanklesaria in #5205
Improve Pytree flatten/unflatten by @copybara-service[bot] in #5216
Pin optax version to fix CI issues by @samanklesaria in #5225
Pin doc deps to lower versions by @samanklesaria in #5222
Fix Sequential unwrapping for static attributes by @samanklesaria in #5218
feat(nnx): add Grouped Query Attention (GQA) support by @ayulockedin in #5180
chore: Migrate gsutil usage to gcloud storage by @gurusai-voleti in #5230
Makes linen.Partitioned compatible with linen.WeightNorm by @copybara-service[bot] in #5234
maintain data/static definition in split / state by @copybara-service[bot] in #5227
Adds a test for linen.WeightNorm and linen.with_partitioning compatibility by @copybara-service[bot] in #5236
Add Tunix link by @copybara-service[bot] in #5239
Add tree-mode support for nnx.shard_map by @copybara-service[bot] in #5238
Rename sharding_names to out_sharding in NNX Variable metadata by @copybara-service[bot] in #5215

New Contributors

@gurusai-voleti made their first contribution in #5230

Full Changelog: v0.12.3...v0.12.4

@vfdev-5

What's Changed

Set numpy<2.4 to fix DeprecationWarning in the CI doctest by @vfdev-5 in #5163
ignore optax deprecation warning by @cgarciae in #5165
fix general guides landing page by @ivrolan in #5139
Remove nnx.split_rngs call wrapping nnx.scan in linen to nnx tutorial by @samanklesaria in #5160
Add out_sharding arguments to linear layers where supported by @jackopenn in #5156
fix kw_only_dataclasses for python 3.14 (part 2) by @copybara-service[bot] in #5135
No public description by @copybara-service[bot] in #5164
Have _graph_flatten respect nnx.data declarations (extension of #5140) by @samanklesaria in #5159
[pmap] Avoid degraded performance under the new jax.pmap. by @copybara-service[bot] in #5152
Support multiple None and UNCONSTRAINED when resolving logical rules by @copybara-service[bot] in #5129
improve hijax guide by @cgarciae in #5115
docs: fix typo 'paramater' -> 'parameter' by @ayulockedin in #5166
Make nnx.pop remove sown attributes. by @samanklesaria in #5133
Rename sharding_names to sharding_metadata by @samanklesaria in #5089
Fix bug in graph overhead benchmark by @samanklesaria in #5183
Fixed typos in the docstrings using antigravity by @vfdev-5 in #5145
docs(nnx): add missing functional args to Conv and LinearGeneral by @ayulockedin in #5174
empty change by @copybara-service[bot] in #5184
Use nnx split during tabulate (clone of #5069) by @samanklesaria in #5186
Handle pure bodies in nnx.fori_loop by @samanklesaria in #5141
Typo fix in _cached_partial method by @vfdev-5 in #5142
Update mnist example to use NNX (clone of #5064) by @samanklesaria in #5188
Docs: Fix typo and clarify introduction in Functional API section by @Moriyuki-S in #5157
Shard split rngs in lifted vmap if spmd_axis_name is given and applied to vmapped axis. Without this, jax.vmap gives an error: ValueError: Mapped away dimension of inputs passed to vmap should be sharded the same. Got inconsistent axis specs: None vs batch due to split_rngs being replicated. by @copybara-service[bot] in #5189
Remove ref cycles introduced by self-calling nested functions. by @copybara-service[bot] in #5193
Add HijaxTransformCoverageTest by @cgarciae in #5190
allow nnx standalone import p1 by @copybara-service[bot] in #5196
Add _graph_node_set_key method for List class by @samanklesaria in #5171
_apply_sharding disallow mixed Explicit/Auto mesh by @copybara-service[bot] in #5199
update flax to version 0.12.3 by @copybara-service[bot] in #5206

New Contributors

@ivrolan made their first contribution in #5139
@jackopenn made their first contribution in #5156
@ayulockedin made their first contribution in #5166
@Moriyuki-S made their first contribution in #5157

Full Changelog: v0.12.2...v0.12.3

0.12.12

What's Changed

Contributors

Uh oh!

0.12.11

What's Changed

New Contributors

Contributors

Uh oh!

0.12.10

What's Changed

Contributors

Uh oh!

Release 0.12.9

What's Changed

New Contributors

Contributors

Uh oh!

0.12.8

What's Changed

New Contributors

Contributors

Uh oh!

0.12.7

What's Changed

New Contributors

Contributors

Uh oh!

0.12.6

What's Changed

Contributors

Uh oh!

0.12.5

What's Changed

Contributors

Uh oh!

0.12.4

What's Changed

New Contributors

Contributors

Uh oh!

0.12.3

What's Changed

New Contributors

Contributors

Uh oh!