google/flax

Pull Requests Commits

add graph_updates argument to jit

Cristian Garcia committed 37 days ago

a1899fd4

Allow DenyList to be compared.

levskaya committed 39 days ago

22468a1b

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

a-googler committed 40 days ago

9b399162

support out_sharding mapping in set_metadata

Cristian Garcia committed 42 days ago

4614ac2b

fix eval_shape's _to_variable

Cristian Garcia committed 43 days ago

8335dfe1

Merge pull request #5247 from samanklesaria:eval_shape_sharding

a-googler committed 43 days ago

49e34d17

Resolve code review comments

samanklesaria committed 43 days ago

4e222b6a

Support explicit sharding for nnx.eval_shape

samanklesaria committed 43 days ago

6bce1438

Change 'sharding' name to 'out_sharding'

samanklesaria committed 43 days ago

3e3c4594

Added sharding propagation support in nnx.eval_shape

vfdev-5 committed 43 days ago

70351571

avoid __eq__ in test_iter_graph

Cristian Garcia committed 44 days ago

39b7613d

add recursive_map example

Cristian Garcia committed 44 days ago

b6a683ea

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

a-googler committed 44 days ago

2970c6ff

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

a-googler committed 44 days ago

ef6e4b17

Merge pull request #5292 from samanklesaria:no_optax_pin

a-googler committed 44 days ago

81a7e42e

add with_attributes

Cristian Garcia committed 44 days ago

bedfdaec

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

jimlinntu committed 45 days ago

3229c96d

simplify set_view API

Cristian Garcia committed 45 days ago

8d7f4863

simplify nnx.display for Variable

Cristian Garcia committed 45 days ago

bba63ce1

update version to 0.12.5

Cristian Garcia committed 45 days ago

6ec76a2b

remove experimental Variable.mutable property

Cristian Garcia committed 45 days ago

fc245892

improve error messages for tree mode errors

Cristian Garcia committed 45 days ago

9c33569d

Merge pull request #5295 from samanklesaria:remove_chex

a-googler committed 45 days ago

bbfab0a6

Merge pull request #5294 from google:nnx-view

a-googler committed 46 days ago

357b2984

expose and update view guide

cgarciae committed 46 days ago

d5ce4368

add error handling to iter_children

Cristian Garcia committed 46 days ago

e067f391

Update tests/jax_utils_test.py

samanklesaria committed 47 days ago

Verified afc169eb

Remove chex dependency

samanklesaria committed 47 days ago

1cdfa4e1

add set_graph_mode

Cristian Garcia committed 47 days ago

972c34e9

No public description

danielsuo committed 47 days ago

5486a636

Newer Older