Search results

nnx.Embed(num_embedding=1) incompatible with "jax_disable_jit" and "debug_nans"

Status: Open.

#4557 In google/flax;

· anthony-Neo opened on Feb 20, 2025

Improve support for optax.GradientTransformationExtraArgs in NNX

Status: Open.

#4545 In google/flax;

· NKlug opened on Feb 13, 2025

AttributeError: module 'jax._src.config' has no attribute 'enable_memories'

Status: Open.

#4535 In google/flax;

· jeromemassot opened on Feb 9, 2025

Restoring a checkpoint with Orbax CheckpointManager fails

Status: Open.

#4521 In google/flax;

· bednarikjan opened on Feb 1, 2025

jax.eval_shape modifies nnx.Module state when using nnx.vmap and nnx.scan

Status: Open.

#4520 In google/flax;

· ytsmiling opened on Feb 1, 2025

[linen] Undocumented RNN errors

Status: Open.

#4513 In google/flax;

· MRiabov opened on Jan 29, 2025

Why does LSTMCell keep rngs in its state?

Status: Open.

#4509 In google/flax;

· JoaoAparicio opened on Jan 28, 2025

How to unroll a model end to end?

Status: Open.

#4507 In google/flax;

· JoaoAparicio opened on Jan 27, 2025

nnx.make_causal_mask() usage

Status: Open.

#4505 In google/flax;

· windmaple opened on Jan 25, 2025

How to use (jax) pytrees inside of nnx modules?

Status: Open.

#4497 In google/flax;

· PhilipVinc opened on Jan 22, 2025

Pool functions reduce over batch dimension and not last dimension

Status: Open.

#4494 In google/flax;

· simonschoelly opened on Jan 21, 2025

Update dependencies for flaxlib

Status: Open.

#4491 In google/flax;

· Bot-wxt1221 opened on Jan 19, 2025