Search results

<code>nnx.cached_partial</code> breaks <code>nnx.Module.sow</code> -&gt; <code>nnx.pop</code> pattern

Status: Open.

#5130 In google/flax;

· am001122 opened on Dec 9, 2025

nnx.data is lost during nnx.eval_shape

Status: Open.

#5128 In google/flax;

· qGentry opened on Dec 8, 2025

Out sharding for modules initialized with JIT is incorrect

Status: Open.

#5127 In google/flax;

· qGentry opened on Dec 8, 2025

Examples how to use fori_loop for gradient accumulation &amp; 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