Commits

Commit history

1. Commits on Nov 10, 2026

   • Add warmup_steps support to nnx.EMA.

     Show description for 204a819
     

     Flax Team

     committed
     204a819View commit details

     Copy full SHA for 204a819

     Browse repository at this point
2. Commits on Nov 7, 2026

   • Fix pre-existing Pyrefly type-check errors in Flax and remove migration suppressions

     Show description for 54aada1
     

     IvyZX

     authored and

     Flax Authors

     committed
     54aada1View commit details

     Copy full SHA for 54aada1

     Browse repository at this point
3. Commits on Nov 5, 2026

   • Align Python 3.12–3.15 classifiers, GitHub Actions CI, and Pyrefly type checking with JAX

     Show description for 71c5b9c
     

     IvyZX

     authored and

     Flax Authors

     committed
     71c5b9cView commit details

     Copy full SHA for 71c5b9c

     Browse repository at this point
4. Commits on Nov 3, 2026

   • Use `jax.tree_util.is_tree_node` instead of private `_registry` in NNX

     Show description for b3243f7
     

     IvyZX

     authored and

     Flax Authors

     committed
     b3243f7View commit details

     Copy full SHA for b3243f7

     Browse repository at this point
5. Commits on Oct 27, 2026

   • Release Flax 0.12.12

     Show description for e937d47
     

     Cristian Garcia

     authored and

     Flax Authors

     committed
     e937d47View commit details

     Copy full SHA for e937d47

     Browse repository at this point

   • rollback of breaking change

     Show description for 62ad9d1
     

     levskaya

     authored and

     Flax Authors

     committed
     62ad9d1View commit details

     Copy full SHA for 62ad9d1

     Browse repository at this point
6. Commits on Oct 26, 2026

   • Merge pull request #5593 from mayuriphad:fix-flatten-dict-sep

     Show description for 6866a65
     

     Flax Authors

     committed
     6866a65View commit details

     Copy full SHA for 6866a65

     Browse repository at this point

   • Merge pull request #5592 from Anu27n:fix/nnx-mha-decode-is-causal

     Show description for 3c3ce3b
     

     Flax Authors

     committed
     3c3ce3bView commit details

     Copy full SHA for 3c3ce3b

     Browse repository at this point

   • Merge pull request #5582 from jdymitarai:fix/nnx-capture-scan-list

     Show description for bf06ba3
     

     Flax Authors

     committed
     bf06ba3View commit details

     Copy full SHA for bf06ba3

     Browse repository at this point

   • Add test for invalid keys with sep

     

     mayuriphad

     authored and

     samanklesaria

     committed
     8a285abView commit details

     Copy full SHA for 8a285ab

     Browse repository at this point

   • Update flatten_dict docstring

     

     mayuriphad

     authored and

     samanklesaria

     committed
     12e226bView commit details

     Copy full SHA for 12e226b

     Browse repository at this point

   • Fix flatten_dict to raise ValueError on invalid keys with sep

     

     mayuriphad

     authored and

     samanklesaria

     committed
     d2bdce7View commit details

     Copy full SHA for d2bdce7

     Browse repository at this point

   • Fix NNX MultiHeadAttention is_causal mask during decode.

     Show description for c3721a2
     

     Anu27n

     authored and

     samanklesaria

     committed
     c3721a2View commit details

     Copy full SHA for c3721a2

     Browse repository at this point

   • Release Flax 0.12.11

     Show description for 899109a
     

     Cristian Garcia

     authored and

     Flax Authors

     committed
     899109aView commit details

     Copy full SHA for 899109a

     Browse repository at this point

   • [nnx] Make `nnx.Sequential` iterable and filter `**kwargs` per layer.

     Show description for 7a24e6a
     

     Cristian Garcia

     authored and

     Flax Authors

     committed
     7a24e6aView commit details

     Copy full SHA for 7a24e6a

     Browse repository at this point
7. Commits on Oct 24, 2026

   • fix(nnx): skip Sequence modules in capture to prevent mixed key types (fixes #5571)

     

     jdymitarai

     committed
     ce5acb2View commit details

     Copy full SHA for ce5acb2

     Browse repository at this point
8. Commits on Oct 21, 2026

   • Merge pull request #5623 from IvyZX:fix-5577-weightnorm

     Show description for 479d43d
     

     Flax Authors

     committed
     479d43dView commit details

     Copy full SHA for 479d43d

     Browse repository at this point

   • Merge pull request #5488 from samanklesaria:multimodal_llm_example

     Show description for 2e2876d
     

     Flax Authors

     committed
     2e2876dView commit details

     Copy full SHA for 2e2876d

     Browse repository at this point
9. Commits on Oct 20, 2026

   • Store nnx.WeightNorm scales as Param so they are trainable

     Show description for ae6d283
     IvyZX

     committed
     ae6d283View commit details

     Copy full SHA for ae6d283

     Browse repository at this point

   • Merge remote-tracking branch 'upstream/main' into multimodal_llm_example

     cgarciae

     committed
     3ee9be7View commit details

     Copy full SHA for 3ee9be7

     Browse repository at this point

   • Merge pull request #5472 from samanklesaria:audio_example

     Show description for 7a610fa
     

     Flax Authors

     committed
     7a610faView commit details

     Copy full SHA for 7a610fa

     Browse repository at this point

   • Print shapes instead of exact numerics in Embed docstrings

     IvyZX

     committed
     c15ac73View commit details

     Copy full SHA for c15ac73

     Browse repository at this point

   • Fix mypy errors with jax 0.12.0

     IvyZX

     committed
     82ca84eView commit details

     Copy full SHA for 82ca84e

     Browse repository at this point

   • Merge remote-tracking branch 'origin/main' into audio_example

     IvyZX

     committed
     c3d172bView commit details

     Copy full SHA for c3d172b

     Browse repository at this point
10. Commits on Oct 12, 2026

    • Add `capture_noop` option to `Module.perturb`.

      Show description for 41deea3
      

      Cristian Garcia

      authored and

      Flax Authors

      committed
      41deea3View commit details

      Copy full SHA for 41deea3

      Browse repository at this point
11. Commits on Oct 8, 2026

    • Don't pass linen_meta_type argument when creating AxisMetadata subclasses.

      Show description for 6c9494b
      

      Flax Team

      committed
      6c9494bView commit details

      Copy full SHA for 6c9494b

      Browse repository at this point
12. Commits on Oct 7, 2026

    • Fix Flax GitHub Actions CI (`pytest` `DeprecationWarning` and `pytype` in `examples/gemma/train.py`).

      Show description for 5e60ced
      

      Cristian Garcia

      authored and

      Flax Authors

      committed
      5e60cedView commit details

      Copy full SHA for 5e60ced

      Browse repository at this point
13. Commits on Oct 6, 2026

    • Replace `# pytype: disable` suppressions with `# pyrefly: ignore`

      Show description for 81b168d
      

      oprypin

      authored and

      Flax Authors

      committed
      81b168dView commit details

      Copy full SHA for 81b168d

      Browse repository at this point
14. Commits on Oct 1, 2026

    • Add `nnx_sow_capture_noop` config flag and `capture_noop` option to `Module.sow`.

      Show description for 6cde97d
      

      Cristian Garcia

      authored and

      Flax Authors

      committed
      6cde97dView commit details

      Copy full SHA for 6cde97d

      Browse repository at this point
15. Commits on Sep 25, 2026

    • Merge pull request #5518 from chenkuanliao:fix-issue-5512

      Show description for 4c44124
      

      Flax Authors

      committed
      4c44124View commit details

      Copy full SHA for 4c44124

      Browse repository at this point

    • Merge pull request #5580 from samanklesaria:fix_mypy_hijax_error

      Show description for 78abd68
      

      Flax Authors

      committed
      78abd68View commit details

      Copy full SHA for 78abd68

      Browse repository at this point

    • Fix mypy hijax error

      samanklesaria

      committed
      ce09c08View commit details

      Copy full SHA for ce09c08

      Browse repository at this point
16. Commits on Sep 24, 2026

    • Release Flax 0.12.10 and use attribute lookup for HiPrim fallback

      Show description for d1f0d69
      

      Cristian Garcia

      authored and

      Flax Authors

      committed
      d1f0d69View commit details

      Copy full SHA for d1f0d69

      Browse repository at this point
17. Commits on Sep 16, 2026

    • suppress pyrefly errors related to JAX scalar types

      Show description for 4f42ed0
      

      Jake VanderPlas

      authored and

      Flax Authors

      committed
      4f42ed0View commit details

      Copy full SHA for 4f42ed0

      Browse repository at this point
18. Commits on Sep 15, 2026

    • Add sepreformer example

      samanklesaria

      committed
      02d3065View commit details

      Copy full SHA for 02d3065

      Browse repository at this point