-
Notifications
You must be signed in to change notification settings - Fork 831
Pull requests: google/flax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
[No functional change] Delete VJPHiPrimitive and replace with HiPrim. This is a simple replace.
#5567
opened Aug 21, 2026 by
copybara-service
Bot
Loading…
fix(traverse_util): validate keys in flatten_dict when separator is specified (#5554)
#5564
opened Aug 19, 2026 by
Ultron09
Loading…
docs(attention): clarify make_causal_mask and make_attention_mask docstrings (#1520)
#5563
opened Aug 19, 2026 by
Ultron09
Loading…
4 tasks done
fix(nnx): make ToLinen hashable for JAX JIT (#4156)
#5562
opened Aug 19, 2026 by
Ultron09
Loading…
4 tasks done
docs(attention): document None default behavior in MultiHeadAttention layers
#5561
opened Aug 19, 2026 by
Ultron09
Loading…
4 tasks done
Add Python 3.14 support and roll up python type check to 3.13 (matching internal).
#5560
opened Aug 19, 2026 by
copybara-service
Bot
Loading…
traverse_util: reject unrestorable keys in flatten_dict with a separator
#5555
opened Aug 15, 2026 by
Amey-Thakur
Loading…
docs(nnx): fill missing docstrings in variables/rnglib, fix helpers g…
#5545
opened Aug 5, 2026 by
greekera1000
Loading…
3 of 4 tasks
docs(nnx): fix training API reference inconsistencies
#5544
opened Aug 4, 2026 by
greekera1000
Loading…
3 of 4 tasks
Enable full NumPy indexing compliance for jax.Ref
#5542
opened Aug 1, 2026 by
copybara-service
Bot
Loading…
Fix vmap axes for input_positions in dot_product_attention_with_rope
#5537
opened Jul 26, 2026 by
TheSaiEaranti
Loading…
3 tasks done
Handle captures in jax transforms with hijax
#5531
opened Jul 21, 2026 by
samanklesaria
Collaborator
Loading…
Preserve RngStream counters for passthrough streams in Rngs.split/broadcast
#5529
opened Jul 17, 2026 by
sohumt123
Loading…
3 of 4 tasks
Fix nnx.Dropout crash when broadcast_dims includes a sharded axis
#5528
opened Jul 17, 2026 by
sohumt123
Loading…
3 of 4 tasks
Fix tracing cache misses in nnx.cond and nnx.switch (#5512)
#5525
opened Jul 6, 2026 by
kapellirohith
Loading…
Add pickle/cloudpickle round-trip tests for Linen modules (#1481)
#5524
opened Jul 6, 2026 by
kapellirohith
Loading…
2 of 4 tasks
Speed up NNX graph flattening by avoiding unnecessary array-ref checks
#5523
opened Jul 6, 2026 by
KookiesNKareem
Loading…
1 of 4 tasks
Update NNX docs for deprecated Variable.value access
#5521
opened Jul 2, 2026 by
sumitdotml
Loading…
2 of 4 tasks
Previous Next
ProTip!
Mix and match filters to narrow down what you’re looking for.