All issues

Search results

Jax and Flax imports out of sync

Status: Open.

#5596 In google/flax;

· JulienRaynal opened on Sep 22, 2026

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

Status: Open.

#5594 In google/flax;

· vineethsaivs opened on Sep 21, 2026

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

Status: Open.

#5591 In google/flax;

· Anu27n opened on Sep 20, 2026

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

Status: Open.

#5589 In google/flax;

· ajitashwath opened on Sep 19, 2026

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

Status: Open.

#5577 In google/flax;

· 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

Status: Open.

#5576 In google/flax;

· PizzasBear opened on Aug 29, 2026

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

Status: Open.

#5571 In google/flax;

· maxencefaldor opened on Aug 26, 2026

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

Status: Open.

#5554 In google/flax;

· Amey-Thakur opened on Aug 15, 2026

Jax-Transform based Sow Implementation

Status: Open.

#5538 In google/flax;

· samanklesaria opened on Jul 28, 2026

flax.nnx.cond causes tracing cache misses

Status: Open.

#5512 In google/flax;

· futscdav opened on Jun 22, 2026

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

Status: Open.

#5487 In google/flax;

· geo-chen opened on Jun 9, 2026

Failed build: CI - with JAX nightly

build failed

Status: Open.

#5473 In google/flax;

· github-actions opened on May 27, 2026