Search results

Feature request, <code>nnx.Sequential()</code> should be the identity function.

Status: Open.

#4782 In google/flax;

· SobhanMP opened on Jun 15, 2025

nnx.scan working

Status: Open.

#4780 In google/flax;

· ak24watch opened on Jun 14, 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

nnx.jit() cannot specify backend or device due to out_shardings

Priority: P1 - soonResponse within 5 business days. Resolution within 30 days. (Assignee required)Response within 5 business days. Resolution within 30 days. (Assignee required)

Status: Open.

#4774 In google/flax;

· wookayin opened on Jun 9, 2025

Scan over layers does not work with jax.grad?

Status: Open.

#4771 In google/flax;

· eirikfagerbakke opened on Jun 5, 2025

JAX and Flax related issue: https://github.com/jax-ml/jax/issues/28937#issue-3083667493

Status: Open.

#4759 In google/flax;

· Beiglar opened on May 22, 2025

Multi-host migration of <code>nnx.pmap</code> to <code>nnx.shard_map</code>

Status: Open.

#4751 In google/flax;

· jaketae opened on May 16, 2025

Return attention weights

Status: Open.

#4744 In google/flax;

· dibgerge opened on May 13, 2025

Improve Gemma example

Status: Open.

#4740 In google/flax;

· cgarciae opened on May 7, 2025

Cannot <code>nnx.eval_shape</code> Modules with <code>BatchNorm</code>

Status: Open.

#4734 In google/flax;

· LarsKue opened on Apr 29, 2025

How to train an NNX netowork using scan and jit?

Status: Open.

#4722 In google/flax;

· cdelv opened on Apr 20, 2025