google/flax

Pull Requests Commits

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

a-googler committed 10 days ago

d48b501f

improve with_rngs

Cristian Garcia committed 14 days ago

983506b3

expose nnx pooling layers in API docs

Cristian Garcia committed 14 days ago

70f4aa96

add Rngs.broadcast

Cristian Garcia committed 14 days ago

3645777b

Cristian Garcia committed 14 days ago

cd619533

Rename pure -> as_pure

samanklesaria committed 14 days ago

61321559

Deprecate the keep_rngs argument in recurrent.py

a-googler committed 15 days ago

3c681943

Rename abstract_with_sharding to as_abstract

samanklesaria committed 15 days ago

dc218b40

Merge pull request #5386 from samanklesaria:fix_doctest

a-googler committed 16 days ago

6d86546f

Merge pull request #5382 from samanklesaria:deprecation_func

a-googler committed 16 days ago

6172e69b

Fix doctests for optimizer update

samanklesaria committed 16 days ago

0d4e25e6

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

a-googler committed 16 days ago

66bed3ab

Add iter_module_children() to iterate over module children only.

hyunn9973 committed 17 days ago

daf3548b

Merge pull request #5383 from samanklesaria:with_rngs

a-googler committed 20 days ago

115cc4d1

Add error message for with_rngs

samanklesaria committed 20 days ago

dd95feea

Add deprecations file

samanklesaria committed 20 days ago

1a9a0683

samanklesaria committed 20 days ago

56527e9b

Add with_rngs function

samanklesaria committed 20 days ago

902b953d

Removed bonsai from docs

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

c8e3d84d

Add nnx.make_jaxpr.

a-googler committed 21 days ago

d539ce78

More Code updates

vfdev-5 committed 21 days ago

5260a1ed

Fixed broken url in the sow warning

vfdev-5 committed 21 days ago

Verified c3e9c743

Added mixed precision: computation and weights dtypes

vfdev-5 committed 22 days ago

cdf3b0c3

Use grain in gemma example

vfdev-5 committed 22 days ago

199e6530

Merge pull request #5353 from samanklesaria:fix_orbax_export_dep

a-googler committed 23 days ago

35d911a2

Merge pull request #5231 from samanklesaria:opt_sharding

a-googler committed 23 days ago

bb50f212

Update pytree guide to use nnx.capture

samanklesaria committed 24 days ago

9f0ea699

fix transform_metadata

Cristian Garcia committed 26 days ago

a138d9f3

add nnx.map and nnx.abstract_with_sharding

Cristian Garcia committed 27 days ago

1572c84e

Newer Older