Search results

NNX: Support for read-only <code>Module</code>s in closures via something like a<code>nnx.freeze</code> or relaxed trace verification

Status: Open.

#5291 In google/flax;

· burnpanck opened on Feb 28, 2026

Installation fails on Windows

Status: Open.

#5260 In google/flax;

· JeppeKlitgaard opened on Feb 19, 2026

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

Status: Open.

#5202 In google/flax;

· divye-joshi opened on Jan 26, 2026

Improvement on the Training API

Status: Open.

#5194 In google/flax;

· vfdev-5 opened on Jan 19, 2026

<code>nnx.Optimizer.update</code> should return the <code>updates</code>

Status: Open.

#5182 In google/flax;

· am001122 opened on Jan 10, 2026

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

Status: Open.

#5176 In google/flax;

· coder0143 opened on Jan 7, 2026

[docs]: Improve docs for NNX API Reference

Status: Open.

#5161 In google/flax;

· ibbyml opened on Jan 3, 2026

Out sharding for modules initialized with JIT is incorrect

Status: Open.

#5127 In google/flax;

· qGentry opened on Dec 8, 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

<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