Issues

• nnx.EMA.update fails when the model holds RNG state (e.g. nnx.Dropout with rngs)

  google/flax#5601 · Filter by author Arthur031221 opened on Sep 27, 2026

• nnx.GRUCell is missing the b_hn bias that its docstring and linen.GRUCell both specify

  google/flax#5594 · Filter by author vineethsaivs opened on Sep 21, 2026

• NNX MultiHeadAttention: is_causal=True during decode attends only to key position 0

  google/flax#5591 · Filter by author Anu27n opened on Sep 20, 2026
  1

• NNXMeta.add_axis / remove_axis are no-ops, silently dropping axis metadata under Linen scan / vmap

  google/flax#5589 · Filter by author ajitashwath opened on Sep 19, 2026

• The scales field of the NNX WeightNorm module isn't learnable.

  google/flax#5577 · Filter by author PizzasBear opened on Aug 29, 2026

• NNX spectral and weight normalisation writes normalised weights back into params diverging from Linen and from the papers' update rules

  google/flax#5576 · Filter by author PizzasBear opened on Aug 29, 2026
  4

• nnx.capture fails with nnx.scan when the model contains nnx.List (mixed str/int pytree dict keys)

  google/flax#5571 · Filter by author maxencefaldor opened on Aug 26, 2026
  1

  

• flatten_dict with a separator cannot round-trip keys containing the separator, and raises an opaque error on non-string keys

  google/flax#5554 · Filter by author Amey-Thakur opened on Aug 15, 2026

• Jax-Transform based Sow Implementation

  google/flax#5538 · Filter by author samanklesaria opened on Jul 28, 2026

• [Security] Arbitrary host-file disclosure when loading an untrusted checkpoint: tensorstore file kvstore follows symlinked chunk files with no containment

  google/flax#5487 · Filter by author geo-chen opened on Jun 9, 2026
  11

• Failed build: CI - with JAX nightly

  build failed
  google/flax#5473 · Filter by author github-actions opened on May 27, 2026
  70

• Add more examples in the NNX docs

  google/flax#5433 · Filter by author vfdev-5 opened on Apr 20, 2026
  4

  

• Installation fails on Windows

  google/flax#5260 · Filter by author JeppeKlitgaard opened on Feb 19, 2026
  6

• feat(nnx): Missing object-oriented pooling layers in NNX

  google/flax#5202 · Filter by author divye-joshi opened on Jan 26, 2026
  2

• Adding transformer encoder and decoder layers to flax source as in pytorch

  google/flax#5176 · Filter by author coder0143 opened on Jan 7, 2026
  3

• [docs]: Improve docs for NNX API Reference

  google/flax#5161 · Filter by author ibbyml opened on Jan 3, 2026
  5

• Out sharding for modules initialized with JIT is incorrect

  google/flax#5127 · Filter by author qGentry opened on Dec 8, 2025
  7

• nnx.Optimizer doesn't respect extra sharding axis added by nnx.scan/nnx.vmap

  Bug

  google/flax#5112 · Filter by author qGentry opened on Nov 24, 2025
  11

• nnx.cached_partial errors on jax Arrays not wrapped in nnx.Variable

  Bug

  google/flax#5109 · Filter by author am001122 opened on Nov 23, 2025
  10

• Outdated and Broken Documentation: Module & Pytree Guide

  google/flax#5100 · Filter by author wtfnukee opened on Nov 18, 2025
  3

• Migrate vae example from flax.linen to flax.nnx

  google/flax#5068 · Filter by author sanepunk opened on Nov 3, 2025
  2

• Feature request: Semantic checkpoint in training loop

  needs-discussion
  google/flax#5059 · Filter by author elly99-AI opened on Oct 29, 2025
  3

• Tensorflow test dependency needs upgrade to support orbax

  google/flax#5054 · Filter by author samanklesaria opened on Oct 27, 2025
  1

• nnx.remat (and other transformations) is easily misused.

  google/flax#5053 · Filter by author liudangyi opened on Oct 25, 2025
  5

  

• Give RngStream the same interface as Jax's stateful RNGs

  google/flax#5046 · Filter by author samanklesaria opened on Oct 23, 2025