Repository Issue Activity (beta)

#5512 flax.nnx.cond causes tracing cache misses

#5433 Add more examples in the NNX docs

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

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

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

#5483 `metrics` does not work well with `vmap`

#5473 Failed build: CI - with JAX nightly

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

#5487 [Security] Arbitrary host-file disclosure when loading an untrusted checkpoint: tensorstore file kvstore follows symlinked chunk files with no containment

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

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

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

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

#5109 `nnx.cached_partial` errors on jax `Array`s not wrapped in `nnx.Variable`

#5495 flax.nnx.max_pool/avg_pool/min_pool raise cryptic errors for wrong-type window_shape/strides/padding

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

#5486 [docs] NNX MultiHeadAttention docs show is_causal, but latest PyPI release does not support it

#5161 [docs]: Improve docs for NNX API Reference

#5468 [NNX] TernaryDense — ternary-weight dense layer for BitNet b1.58 style models

#5452 custom_vjp fails with `graph_updates=False`.

#4488 Failed build: CI - with JAX nightly

#5462 Proposal: add DoRA (Weight-Decomposed Low-Rank Adaptation) to nnx.nn.lora

#5456 Adding Dynamic Tanh (DyT) element wise Operation (Paper from Meta)

#5457 Adding Dynamic erf (Derf) pointwise function to norms, successor to Dynamic tanh

#4247 Pickling a trained model (NNX)