Search results

flax.training.orbax_utils.restore_args_from_target raises AttributeError on SingleDeviceSharding with the deprecated mesh argument

Status: Open.

#5500 In google/flax;

· gdrosa opened on Jun 12, 2026

flax.training.common_utils.shard raises ZeroDivisionError for a zero-sized array (shape (0, 0))

Status: Open.

#5499 In google/flax;

· gdrosa opened on Jun 12, 2026

flax.training.common_utils.get_metrics raises 'iteration over a 0-d array' for single-step (N=1) pmap output

Status: Open.

#5498 In google/flax;

· gdrosa opened on Jun 12, 2026

flax.nnx.dot_product_attention does not validate dropout_rate (negative/NaN silently disables dropout)

Status: Open.

#5497 In google/flax;

· gdrosa opened on Jun 12, 2026

flax.nnx.dot_product_attention silently returns wrong-shaped output under python -O (rank asserts dropped)

Status: Open.

#5496 In google/flax;

· gdrosa opened on Jun 12, 2026

flax.nnx.merge raises a misleading 'Incorrect number of leaves' ValueError for a wrong-type state argument

Status: Open.

#5494 In google/flax;

· gdrosa opened on Jun 12, 2026

flax.nnx promote_dtype raises TypeError on a single None or 0-d scalar (signature contradicts its *args docstring)

Status: Open.

#5493 In google/flax;

· gdrosa opened on Jun 12, 2026

nnx.MultiMetric triggers JIT error when splitting the model into graph/state

Status: Open.

#5489 In google/flax;

· Nogarx opened on Jun 9, 2026

[Security] Path traversal / arbitrary file read on restore_checkpoint via a malicious multiprocess-array placeholder leaf

Status: Open.

#5487 In google/flax;

· geo-chen opened on Jun 9, 2026

metrics does not work well with vmap

Status: Open.

#5483 In google/flax;

· narang99 opened on Jun 4, 2026

Failed build: CI - with JAX nightly

build failed

Status: Open.

#5473 In google/flax;

· github-actions opened on May 27, 2026

nnx.clone creates buffer copies with the same IDs, causing errors with donate_argnums

Status: Open.

#5461 In google/flax;

· johnlyzhou opened on May 11, 2026