google/flax

Pull Requests Commits

"From PyTorch to JAX and Flax" migration guide

vfdev-5 committed 6 hours ago

fbc5ddc5

Avoid deprecated pattern of constructing a dtype from an array.

Jake VanderPlas committed 3 days ago

9534e266

Merge pull request #5400 from google:fix-ci-2

a-googler committed 5 days ago

a0cc8837

cgarciae committed 5 days ago

a3bd2ca9

change defaults for Dropout and BatchNorm

Cristian Garcia committed 5 days ago

84da828a

Merge pull request #5389 from samanklesaria:rename_pure

a-googler committed 5 days ago

d874278f

Fix the wrong names for the hooks in nnx.with_metadata, fixed such that names match the proper hook names in Variable class.

a-googler committed 5 days ago

5a4859d2

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.

danielsuo committed 6 days ago

355e141a

Merge pull request #5380 from google:removed-bonsai-from-docs

a-googler committed 7 days ago

d48b501f

improve with_rngs

Cristian Garcia committed 10 days ago

983506b3

expose nnx pooling layers in API docs

Cristian Garcia committed 10 days ago

70f4aa96

add Rngs.broadcast

Cristian Garcia committed 10 days ago

3645777b

Cristian Garcia committed 10 days ago

cd619533

Rename pure -> as_pure

samanklesaria committed 11 days ago

61321559

Deprecate the keep_rngs argument in recurrent.py

a-googler committed 11 days ago

3c681943

Merge pull request #5386 from samanklesaria:fix_doctest

a-googler committed 12 days ago

6d86546f

Merge pull request #5382 from samanklesaria:deprecation_func

a-googler committed 12 days ago

6172e69b

Fix doctests for optimizer update

samanklesaria committed 12 days ago

0d4e25e6

Merge pull request #5214 from anshuljoshi:optimizer-return-updates

a-googler committed 12 days ago

66bed3ab

Add iter_module_children() to iterate over module children only.

hyunn9973 committed 13 days ago

daf3548b

Merge pull request #5383 from samanklesaria:with_rngs

a-googler committed 16 days ago

115cc4d1

Add error message for with_rngs

samanklesaria committed 16 days ago

dd95feea

Add deprecations file

samanklesaria committed 16 days ago

1a9a0683

samanklesaria committed 16 days ago

56527e9b

Add with_rngs function

samanklesaria committed 16 days ago

902b953d

Removed bonsai from docs

vfdev-5 committed 16 days ago

Verified 86546581

Replaced custom PriorityStr with standard tuple sorting in _graph_node_flatten. Python 3.13.11 optimized str comparison bypasses custom subclasses (causing TypeError), and used getattr in _setattr because unflattening order becomes non-deterministic (preventing AttributeError).

a-googler committed 17 days ago

c8e3d84d

Add nnx.make_jaxpr.

a-googler committed 17 days ago

d539ce78

Merge pull request #5353 from samanklesaria:fix_orbax_export_dep

a-googler committed 19 days ago

35d911a2

Merge pull request #5231 from samanklesaria:opt_sharding

a-googler committed 19 days ago

bb50f212

Older