Flax Files

What's Changed

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

New Contributors

• @dparikh79 made their first contribution in https://github.com/google/flax/pull/5475

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