google/flax

Pull Requests Commits

Added sharding propagation support in nnx.eval_shape

vfdev-5 committed 51 days ago

47626d88

Merge pull request #5095 from google:jax-minver-0.8.1

a-googler committed 56 days ago

697f4e5c

update jax minver to 0.8.1

cgarciae committed 56 days ago

b424c49e

Merge pull request #5097 from google:fix-main-4

a-googler committed 57 days ago

bee81ed4

Fixed deprecated .value usage failing CI tests

vfdev-5 committed 57 days ago

00cf44e5

Merge pull request #5084 from samanklesaria:variable_eq

a-googler committed 58 days ago

09537f49

Merge pull request #5052 from google:deprecate-value-2

a-googler committed 58 days ago

b7e66b74

deprecate .value

cgarciae committed 58 days ago

a2eac321

Merge pull request #5055 from mohsinm-dev:fix-bound-method-auto-unbinding

a-googler committed 58 days ago

d6e8e14f

Merge pull request #5073 from google:fix-nnx-scan-carry-pytree-sow

a-googler committed 58 days ago

544a0505

Merge pull request #5071 from chapman20j:set_mode_info

a-googler committed 58 days ago

0b1d841c

Merge branch 'main' into pr/5055

samanklesaria committed 58 days ago

dece6b34

Merge pull request #5090 from google:fix-main-3

a-googler committed 61 days ago

5109e2c5

cgarciae committed 61 days ago

d982784d

backfill eager_sharding in set_metadata

Cristian Garcia committed 61 days ago

1ac99b09

add HijaxVariable [rollforward]

Cristian Garcia committed 61 days ago

b5db513f

Add __eq__ for variables

samanklesaria committed 62 days ago

1d2a672b

Merge pull request #5079 from samanklesaria:eager_sharding_context

a-googler committed 63 days ago

7710c304

Merge pull request #5074 from google:params_metadata_args

a-googler committed 63 days ago

b08cb206

replaced carry_graphdefs -> graphdefs

vfdev-5 committed 63 days ago

aa02f704

Implement *_metadata to other layers

vfdev-5 committed 63 days ago

f320a5c2

Added kernel_metadata/bias_metadata args to nnx layers

vfdev-5 committed 63 days ago

c6036553

Merge pull request #4990 from samanklesaria:issues/4593

a-googler committed 63 days ago

55a23669

Merge pull request #5066 from google:dataclass

a-googler committed 63 days ago

74b72bd5

add nnx dataclass

cgarciae committed 63 days ago

5b4c600a

Merge pull request #5080 from google:improve-promote-dtype-support

a-googler committed 64 days ago

49ef97c3

Merge pull request #5083 from dan-zheng:docs

a-googler committed 64 days ago

b1f573b7

Update gspmd guide to mention Variable init with eager_sharding

samanklesaria committed 64 days ago

719f6879

Merge branch 'main' into eager_sharding_context

samanklesaria committed 64 days ago

0f5cdf14

Add nnx.use_eager_sharding context manager

samanklesaria committed 64 days ago

f600dcb5

Older