google/flax

Pull Requests Commits

forward fix simple transforms by default

Cristian Garcia committed 4 days ago

6bffabbb

expose nnx pooling layers in API docs

Cristian Garcia committed 4 days ago

70f4aa96

add Rngs.broadcast

Cristian Garcia committed 4 days ago

3645777b

Cristian Garcia committed 4 days ago

cd619533

Deprecate the keep_rngs argument in recurrent.py

a-googler committed 5 days ago

3c681943

Merge pull request #5386 from samanklesaria:fix_doctest

a-googler committed 6 days ago

6d86546f

Merge pull request #5382 from samanklesaria:deprecation_func

a-googler committed 6 days ago

6172e69b

Fix doctests for optimizer update

samanklesaria committed 6 days ago

0d4e25e6

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

a-googler committed 7 days ago

66bed3ab

Add iter_module_children() to iterate over module children only.

hyunn9973 committed 7 days ago

daf3548b

Merge pull request #5383 from samanklesaria:with_rngs

a-googler committed 10 days ago

115cc4d1

Add error message for with_rngs

samanklesaria committed 10 days ago

dd95feea

Add deprecations file

samanklesaria committed 10 days ago

1a9a0683

samanklesaria committed 10 days ago

56527e9b

Add with_rngs function

samanklesaria committed 10 days ago

902b953d

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 11 days ago

c8e3d84d

Add nnx.make_jaxpr.

a-googler committed 12 days ago

d539ce78

Merge pull request #5353 from samanklesaria:fix_orbax_export_dep

a-googler committed 13 days ago

35d911a2

Merge pull request #5231 from samanklesaria:opt_sharding

a-googler committed 13 days ago

bb50f212

fix transform_metadata

Cristian Garcia committed 17 days ago

a138d9f3

add nnx.map and nnx.abstract_with_sharding

Cristian Garcia committed 17 days ago

1572c84e

Merge pull request #5246 from samanklesaria:generalize_out_sharding

a-googler committed 17 days ago

3bcc4a91

improve aliasing error msg

Cristian Garcia committed 18 days ago

1a567002

add graph node in prefix checks

Cristian Garcia committed 18 days ago

9eb565cd

improve nnx.Dict error handling

Cristian Garcia committed 18 days ago

ebc81850

simplify SimpleScan

Cristian Garcia committed 18 days ago

2965f95a

Do a few more cleanups after the pmap shmap merge deletion.

yashk2810 committed 18 days ago

f521e955

Do a few more cleanups after pmap_shmap merge

hawkinsp committed 18 days ago

50433ef0

update to version 0.12.6

Cristian Garcia committed 18 days ago

ee6d3c58

Do a few more cleanups after pmap_shmap merge

yashk2810 committed 18 days ago

b96a07a0

Older