google/flax

Pull Requests Commits

use State instead of GraphState

Cristian Garcia committed 1 day ago

d40e1578

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

a-googler committed 2 days ago

a0cc8837

cgarciae committed 2 days ago

a3bd2ca9

change defaults for Dropout and BatchNorm

Cristian Garcia committed 2 days ago

84da828a

Merge pull request #5389 from samanklesaria:rename_pure

a-googler committed 2 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 2 days ago

5a4859d2

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

danielsuo committed 2 days ago

355e141a

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

a-googler committed 3 days ago

d48b501f

improve with_rngs

Cristian Garcia committed 7 days ago

983506b3

expose nnx pooling layers in API docs

Cristian Garcia committed 7 days ago

70f4aa96

add Rngs.broadcast

Cristian Garcia committed 7 days ago

3645777b

Cristian Garcia committed 7 days ago

cd619533

Rename pure -> as_pure

samanklesaria committed 7 days ago

61321559

Deprecate the keep_rngs argument in recurrent.py

a-googler committed 8 days ago

3c681943

Merge pull request #5386 from samanklesaria:fix_doctest

a-googler committed 9 days ago

6d86546f

Merge pull request #5382 from samanklesaria:deprecation_func

a-googler committed 9 days ago

6172e69b

Fix doctests for optimizer update

samanklesaria committed 9 days ago

0d4e25e6

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

a-googler committed 9 days ago

66bed3ab

Add iter_module_children() to iterate over module children only.

hyunn9973 committed 10 days ago

daf3548b

Merge pull request #5383 from samanklesaria:with_rngs

a-googler committed 13 days ago

115cc4d1

Add error message for with_rngs

samanklesaria committed 13 days ago

dd95feea

Add deprecations file

samanklesaria committed 13 days ago

1a9a0683

samanklesaria committed 13 days ago

56527e9b

Add with_rngs function

samanklesaria committed 13 days ago

902b953d

Removed bonsai from docs

vfdev-5 committed 13 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 13 days ago

c8e3d84d

Add nnx.make_jaxpr.

a-googler committed 14 days ago

d539ce78

Merge pull request #5353 from samanklesaria:fix_orbax_export_dep

a-googler committed 16 days ago

35d911a2

Merge pull request #5231 from samanklesaria:opt_sharding

a-googler committed 16 days ago

bb50f212

fix transform_metadata

Cristian Garcia committed 19 days ago

a138d9f3

Older