google/flax

Pull Requests Commits

add tree-mode-nnx FLIP

Cristian Garcia committed 20 hours ago

68766acc

Merge pull request #5226 from google:nnx-optimizer-removed-redund

a-googler committed 20 hours ago

ef6e4b17

Merge pull request #5292 from samanklesaria:no_optax_pin

a-googler committed 20 hours ago

81a7e42e

add with_attributes

Cristian Garcia committed 20 hours ago

bedfdaec

Annotate `wrapped_fn`'s `self` argument in `decorator_lift_transform` for trace-ability

jimlinntu committed 21 hours ago

3229c96d

simplify set_view API

Cristian Garcia committed 1 day ago

8d7f4863

simplify nnx.display for Variable

Cristian Garcia committed 1 day ago

bba63ce1

update version to 0.12.5

Cristian Garcia committed 1 day ago

6ec76a2b

remove experimental Variable.mutable property

Cristian Garcia committed 1 day ago

fc245892

improve error messages for tree mode errors

Cristian Garcia committed 1 day ago

9c33569d

Merge pull request #5295 from samanklesaria:remove_chex

a-googler committed 1 day ago

bbfab0a6

Merge pull request #5294 from google:nnx-view

a-googler committed 2 days ago

357b2984

expose and update view guide

cgarciae committed 2 days ago

d5ce4368

add error handling to iter_children

Cristian Garcia committed 2 days ago

e067f391

Update tests/jax_utils_test.py

samanklesaria committed 2 days ago

Verified afc169eb

Remove chex dependency

samanklesaria committed 2 days ago

1cdfa4e1

add set_graph_mode

Cristian Garcia committed 2 days ago

972c34e9

No public description

danielsuo committed 3 days ago

5486a636

Remove optax pin

samanklesaria committed 3 days ago

48e0f98a

rename graph.py -> graphlib.py

Cristian Garcia committed 3 days ago

c860beef

fix nn.vmap's split_fn

Cristian Garcia committed 6 days ago

babce880

fix nn.vmap's _broadcast_prefix_tree

Cristian Garcia committed 6 days ago

437bb82b

Use linen_vars_to_nnx_attrs in ToLinen variable restoration

Cristian Garcia committed 7 days ago

777b50a8

Add tree-mode-only nnx.{vjp,jvp,jit_partial}

Cristian Garcia committed 8 days ago

9a1bc4eb

Use linen_vars_to_nnx_attrs in ToLinen variable restoration

Cristian Garcia committed 8 days ago

5bcb1f56

Add tree-mode support for nnx.recursive_map, nnx.view, and nnx.view_info

Cristian Garcia committed 12 days ago

3ae65afe

Add tree-mode support to iter_graph

Cristian Garcia committed 12 days ago

8f8ec30a

add tree-mode support for pmap

Cristian Garcia committed 12 days ago

2f6b04c9

Add tree-mode support to nnx.{while,fori}_loop

Cristian Garcia committed 12 days ago

e108ab89

Add tree-mode support for nnx.custom_vjp

Cristian Garcia committed 12 days ago

c9252176

Older