google/flax

Pull Requests Commits

use State instead of GraphState

Cristian Garcia committed 145 days ago

d40e1578

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

a-googler committed 146 days ago

a0cc8837

cgarciae committed 146 days ago

a3bd2ca9

change defaults for Dropout and BatchNorm

Cristian Garcia committed 146 days ago

84da828a

Merge pull request #5389 from samanklesaria:rename_pure

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

5a4859d2

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

danielsuo committed 147 days ago

355e141a

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

a-googler committed 148 days ago

d48b501f

improve with_rngs

Cristian Garcia committed 151 days ago

983506b3

expose nnx pooling layers in API docs

Cristian Garcia committed 151 days ago

70f4aa96

add Rngs.broadcast

Cristian Garcia committed 151 days ago

3645777b

Cristian Garcia committed 151 days ago

cd619533

Rename pure -> as_pure

samanklesaria committed 152 days ago

61321559

Deprecate the keep_rngs argument in recurrent.py

a-googler committed 152 days ago

3c681943

Merge pull request #5386 from samanklesaria:fix_doctest

a-googler committed 153 days ago

6d86546f

Merge pull request #5382 from samanklesaria:deprecation_func

a-googler committed 153 days ago

6172e69b

Fix doctests for optimizer update

samanklesaria committed 153 days ago

0d4e25e6

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

a-googler committed 153 days ago

66bed3ab

Add iter_module_children() to iterate over module children only.

hyunn9973 committed 154 days ago

daf3548b

Merge pull request #5383 from samanklesaria:with_rngs

a-googler committed 157 days ago

115cc4d1

Add error message for with_rngs

samanklesaria committed 157 days ago

dd95feea

Add deprecations file

samanklesaria committed 157 days ago

1a9a0683

samanklesaria committed 157 days ago

56527e9b

Add with_rngs function

samanklesaria committed 157 days ago

902b953d

Removed bonsai from docs

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

c8e3d84d

Add nnx.make_jaxpr.

a-googler committed 158 days ago

d539ce78

Merge pull request #5353 from samanklesaria:fix_orbax_export_dep

a-googler committed 160 days ago

35d911a2

Merge pull request #5231 from samanklesaria:opt_sharding

a-googler committed 160 days ago

bb50f212

fix transform_metadata

Cristian Garcia committed 164 days ago

a138d9f3

Older