Search results

NNX incorrectly handles empty mesh/logical axis rules

Status: Open.

#5124 In google/flax;

· qGentry opened on Dec 4, 2025

Train step with nnx.jit/jax.jit+nnx.grad consumes twice memory than jax.jit+jax.grad equivalent

Status: Open.

#5116 In google/flax;

· qGentry opened on Nov 26, 2025

Examples how to use fori_loop for gradient accumulation & more clear exceptions

Status: Open.

Bug

#5113 In google/flax;

· qGentry opened on Nov 24, 2025

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

Status: Open.

Bug

#5112 In google/flax;

· qGentry opened on Nov 24, 2025

nnx.eval_shape should propagate 'sharding' and 'format' to ShapeDtypeStructs

Status: Open.

Feature

#5110 In google/flax;

· qGentry opened on Nov 24, 2025

<code>nnx.cached_partial</code> errors on jax <code>Array</code>s not wrapped in <code>nnx.Variable</code>

Status: Open.

Bug

#5109 In google/flax;

· am001122 opened on Nov 23, 2025

<code>flax.jax_utils</code> not discoverable from the NNX docs.

Status: Open.

#5107 In google/flax;

· ibbyml opened on Nov 21, 2025

Outdated and Broken Documentation: Module &amp; Pytree Guide

Status: Open.

#5100 In google/flax;

· wtfnukee opened on Nov 18, 2025

[CI] Reenable example tests

Status: Open.

#5099 In google/flax;

· vfdev-5 opened on Nov 18, 2025

Migrate vae example from flax.linen to flax.nnx

Status: Open.

#5068 In google/flax;

· sanepunk opened on Nov 3, 2025

nnx.tabulate expects concrete values in unreleased flax&gt;0.12.0

Status: Open.

Bug

#5067 In google/flax;

· DBraun opened on Nov 3, 2025

Feature request: Semantic checkpoint in training loop

needs-discussion

Status: Open.

#5059 In google/flax;

· elly99-AI opened on Oct 29, 2025