google/flax

Pull Requests Commits

Merge pull request #5292 from samanklesaria:no_optax_pin

a-googler committed 35 days ago

81a7e42e

add with_attributes

Cristian Garcia committed 35 days ago

bedfdaec

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

jimlinntu committed 35 days ago

3229c96d

simplify set_view API

Cristian Garcia committed 36 days ago

8d7f4863

simplify nnx.display for Variable

Cristian Garcia committed 36 days ago

bba63ce1

update version to 0.12.5

Cristian Garcia committed 36 days ago

6ec76a2b

remove experimental Variable.mutable property

Cristian Garcia committed 36 days ago

fc245892

improve error messages for tree mode errors

Cristian Garcia committed 36 days ago

9c33569d

Merge pull request #5295 from samanklesaria:remove_chex

a-googler committed 36 days ago

bbfab0a6

Merge pull request #5294 from google:nnx-view

a-googler committed 36 days ago

357b2984

expose and update view guide

cgarciae committed 36 days ago

d5ce4368

add error handling to iter_children

Cristian Garcia committed 37 days ago

e067f391

Update tests/jax_utils_test.py

samanklesaria committed 37 days ago

Verified afc169eb

Remove chex dependency

samanklesaria committed 37 days ago

1cdfa4e1

add set_graph_mode

Cristian Garcia committed 37 days ago

972c34e9

No public description

danielsuo committed 37 days ago

5486a636

Remove optax pin

samanklesaria committed 37 days ago

48e0f98a

rename graph.py -> graphlib.py

Cristian Garcia committed 37 days ago

c860beef

fix nn.vmap's split_fn

Cristian Garcia committed 41 days ago

babce880

fix nn.vmap's _broadcast_prefix_tree

Cristian Garcia committed 41 days ago

437bb82b

Use linen_vars_to_nnx_attrs in ToLinen variable restoration

Cristian Garcia committed 42 days ago

777b50a8

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

Cristian Garcia committed 43 days ago

9a1bc4eb

Use linen_vars_to_nnx_attrs in ToLinen variable restoration

Cristian Garcia committed 43 days ago

5bcb1f56

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

Cristian Garcia committed 47 days ago

3ae65afe

Add tree-mode support to iter_graph

Cristian Garcia committed 47 days ago

8f8ec30a

add tree-mode support for pmap

Cristian Garcia committed 47 days ago

2f6b04c9

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

Cristian Garcia committed 47 days ago

e108ab89

Add tree-mode support for nnx.custom_vjp

Cristian Garcia committed 47 days ago

c9252176

Add tree-mode support to nnx.scan

Cristian Garcia committed 47 days ago

934d53cd

Add tree-mode support to nnx.{cond,switch,eval_shape,checkify}

Cristian Garcia committed 47 days ago

761a27d3

Newer Older