Search results

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

Make functional versions of <code>iter_modules</code> and <code>iter_children</code> and deprecate them as <code>nnx.Module</code> methods

Status: Open.

#4943 In google/flax;

· vfdev-5 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

Edge case in nnx.tabulate when nnx.Module stores empty dictionary

Status: Open.

#4889 In google/flax;

· DBraun opened on Aug 14, 2025

compute_flops and compute_vjp_flops for nnx.tabulate

Status: Open.

#4883 In google/flax;

· DBraun opened on Aug 10, 2025

Undocumented removal of nnx.Dict in flax 0.11

Status: Open.

#4879 In google/flax;

· jjyyxx opened on Aug 9, 2025

Silent memory explosion related to nnx.scan and closures

awaiting response

Status: Open.

#4871 In google/flax;

· kvablack opened on Aug 5, 2025

Unnecessary check for trace level in nnx.Variable.__setattr__

Status: Open.

#4847 In google/flax;

· liudangyi opened on Jul 25, 2025

<code>nnx.while_loop</code> does not work with <code>Scan over layers</code> or <code>nnx.remat</code>

Status: Open.

#4819 In google/flax;

· gliese581gg opened on Jul 16, 2025

Better type Module.apply signature - PR discussion

Status: Open.

#4779 In google/flax;

· Daraan opened on Jun 13, 2025

Deprecation of <code>nnx.Remat</code> and usage of <code>nnx.remat</code>

Status: Open.

#4778 In google/flax;

· lraedel opened on Jun 11, 2025