Search results

Give RngStream the same interface as Jax's stateful RNGs

Status: Open.

#5046 In google/flax;

· samanklesaria opened on Oct 23, 2025

ValueError when <code>nnx.jit</code> wraps <code>nnx.custom_vjp</code>.

Status: Open.

#5044 In google/flax;

· samanklesaria opened on Oct 22, 2025

Provide a supported way to convert str keys from orbax checkpoints back to int keys

awaiting response

Status: Open.

#5030 In google/flax;

· pcmoritz opened on Oct 15, 2025

Python 3.14 compatibility issue

Status: Open.

#5027 In google/flax;

· burnpanck opened on Oct 14, 2025

Port pytorch migration guide from linen to nnx

Status: Open.

#4997 In google/flax;

· samanklesaria opened on Oct 4, 2025

How to construct freezable Parameters?

needs-discussion

Status: Open.

#4995 In google/flax;

· pfackeldey opened on Oct 3, 2025

NNX JIT fails with static args and in shardings

Status: Open.

Bug

#4989 In google/flax;

· m-braganca opened on Sep 29, 2025

Feature/Docs request: Saving and Loading checkpoints with optimizers and metrics

needs-discussion

Status: Open.

#4979 In google/flax;

· cdelv opened on Sep 24, 2025

Emphasize Explicit Randomness in Documentation

Status: Open.

#4971 In google/flax;

· samanklesaria opened on Sep 22, 2025

Updates on flax, jax and optax caused regressions in google-research projects

needs-discussion

Status: Open.

#4944 In google/flax;

· xvdp opened on Sep 15, 2025

Intermediate value capture API via JAX's <code>hijax.Box</code>

Status: Open.

#4924 In google/flax;

· IvyZX opened on Sep 4, 2025

Feature Request: Automatic Mixed Precision

needs-discussion

Status: Open.

#4921 In google/flax;

· Artoriuz opened on Aug 31, 2025