google/flax

Pull Requests Commits

add graph_updates to while_loop

Cristian Garcia committed 9 hours ago

a8ecd069

add graph_updates for scan

Cristian Garcia committed 14 hours ago

56471a47

Merge pull request #5270 from samanklesaria:split_method

a-googler committed 16 hours ago

d4653e55

add graph_updates to vmap

Cristian Garcia committed 16 hours ago

35972c24

Remove jax/tools/colab_tpu.py.

hawkinsp committed 17 hours ago

8ed67b85

add graph_updates argument in shard_map

Cristian Garcia committed 20 hours ago

115dcf58

Add split method to RngStream

samanklesaria committed 21 hours ago

696328ab

add graph_updates argument to jit

Cristian Garcia committed 1 day ago

a1899fd4

Allow DenyList to be compared.

levskaya committed 3 days ago

22468a1b

Merge pull request #5259 from ayulockedin:feat/gqa-support-mha

a-googler committed 3 days ago

9b399162

support out_sharding mapping in set_metadata

Cristian Garcia committed 6 days ago

4614ac2b

fix eval_shape's _to_variable

Cristian Garcia committed 6 days ago

8335dfe1

Merge pull request #5247 from samanklesaria:eval_shape_sharding

a-googler committed 7 days ago

49e34d17

Resolve code review comments

samanklesaria committed 7 days ago

4e222b6a

Support explicit sharding for nnx.eval_shape

samanklesaria committed 7 days ago

6bce1438

Change 'sharding' name to 'out_sharding'

samanklesaria committed 7 days ago

3e3c4594

Added sharding propagation support in nnx.eval_shape

vfdev-5 committed 7 days ago

70351571

avoid __eq__ in test_iter_graph

Cristian Garcia committed 8 days ago

39b7613d

add recursive_map example

Cristian Garcia committed 8 days ago

b6a683ea

Merge pull request #5237 from mohsinm-dev:fix/pytreenode-generic-parameters

a-googler committed 8 days ago

2970c6ff

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

a-googler committed 8 days ago

ef6e4b17

Merge pull request #5292 from samanklesaria:no_optax_pin

a-googler committed 8 days ago

81a7e42e

add with_attributes

Cristian Garcia committed 8 days ago

bedfdaec

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

jimlinntu committed 8 days ago

3229c96d

simplify set_view API

Cristian Garcia committed 9 days ago

8d7f4863

simplify nnx.display for Variable

Cristian Garcia committed 9 days ago

bba63ce1

update version to 0.12.5

Cristian Garcia committed 9 days ago

6ec76a2b

remove experimental Variable.mutable property

Cristian Garcia committed 9 days ago

fc245892

improve error messages for tree mode errors

Cristian Garcia committed 9 days ago

9c33569d

Merge pull request #5295 from samanklesaria:remove_chex

a-googler committed 9 days ago

bbfab0a6

Older