Go
Home
Pricing
FAQ
Install
Home
Pricing
FAQ
Install
Login
via GitHub
google/flax
Pull Requests
Commits
Open
Closed
Rename abstract_with_sharding to as_abstract
#5388 opened 2026-04-01 16:42 by
samanklesaria
add nnx.prefix
#5387 opened 2026-04-01 01:03 by
copybara-service[bot]
Deprecate the keep_rngs argument in recurrent.py
#5385 opened 2026-03-31 18:50 by
copybara-service[bot]
Rename vars_as to with_vars
#5384 opened 2026-03-30 15:17 by
samanklesaria
Add flag for variables in nnx.merge
#5381 opened 2026-03-27 20:55 by
samanklesaria
Removed bonsai from docs
#5380 opened 2026-03-27 16:16 by
vfdev-5
fix Rngs.__getattr__ returning default stream for dunder names (#5375)
#5379 opened 2026-03-27 00:57 by
mohsinm-dev
Fixed broken url in the sow warning
#5377 opened 2026-03-26 13:52 by
vfdev-5
Consolidate "view" functions with similar naming convention.
#5376 opened 2026-03-25 21:09 by
samanklesaria
[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.
#5374 opened 2026-03-25 09:26 by
copybara-service[bot]
Upgrade third_party/python_runtime/v3_13_unstable to 3.13.11
#5371 opened 2026-03-24 09:54 by
copybara-service[bot]
Update pytree guide to use nnx.capture
#5370 opened 2026-03-23 15:34 by
samanklesaria
Update the RNG guide with split and without StateAxes
#5363 opened 2026-03-20 17:06 by
samanklesaria
Remove spmd_lowering from MeshComputation and do more cleanups after pmap, PmapSharding is gone
#5360 opened 2026-03-20 06:09 by
copybara-service[bot]
improve scaling up guide
#5354 opened 2026-03-18 19:30 by
cgarciae
Rename `SDS(vma: frozenset = ...)` to `SDS(manual_type: jax.sharding.ManualAxisType = ...)`
#5351 opened 2026-03-18 02:16 by
copybara-service[bot]
Add 'prefix' argument to temp_flip_flag function
#5343 opened 2026-03-17 20:28 by
samanklesaria
Fix broken error-message formatting in flax.nnx
#5334 opened 2026-03-16 13:35 by
gdrosa
Add nnx.prefix
#5325 opened 2026-03-12 16:14 by
samanklesaria
add transforms_tree guide
#5307 opened 2026-03-04 18:22 by
cgarciae
Bump the pip group across 9 directories with 4 updates
dependencies
python
#5300 opened 2026-03-04 00:07 by
dependabot[bot]
change defaults for Dropout and BatchNorm
#5299 opened 2026-03-03 14:49 by
copybara-service[bot]
add nnx_basics_tree.ipynb
#5293 opened 2026-03-02 19:16 by
cgarciae
Expand Einsum test coverage with parameterized tests
#5285 opened 2026-02-26 21:35 by
us
Add tests, docstrings and validation to metrics.py
#5284 opened 2026-02-26 21:33 by
us
Add examples and fix docstrings in nnx.nn.dtypes
#5283 opened 2026-02-26 21:11 by
us
Expand Dropout test coverage and migrate to absltest
#5281 opened 2026-02-26 20:32 by
us
Expand LoRA test coverage with dtype, gradient, and edge case tests
#5280 opened 2026-02-26 20:25 by
us
Expand TestLinenConsistency parameterized.product in attention_test
#5279 opened 2026-02-26 20:16 by
us
Expand Embed test coverage with functional, attend, and dtype tests
#5278 opened 2026-02-26 06:26 by
us
Older