google/flax

Pull Requests Commits

Added workaround by P.Hawkins

vfdev-5 committed 47 days ago

67f596e8

Support for python 3.14

vfdev-5 committed 59 days ago

519fb53b

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

a-googler committed 60 days ago

bee81ed4

Fixed deprecated .value usage failing CI tests

vfdev-5 committed 60 days ago

00cf44e5

Merge pull request #5084 from samanklesaria:variable_eq

a-googler committed 61 days ago

09537f49

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

a-googler committed 61 days ago

b7e66b74

deprecate .value

cgarciae committed 61 days ago

a2eac321

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

a-googler committed 61 days ago

d6e8e14f

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

a-googler committed 61 days ago

544a0505

Merge pull request #5071 from chapman20j:set_mode_info

a-googler committed 61 days ago

0b1d841c

Merge branch 'main' into pr/5055

samanklesaria committed 61 days ago

dece6b34

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

a-googler committed 64 days ago

5109e2c5

cgarciae committed 64 days ago

d982784d

backfill eager_sharding in set_metadata

Cristian Garcia committed 64 days ago

1ac99b09

add HijaxVariable [rollforward]

Cristian Garcia committed 64 days ago

b5db513f

Add __eq__ for variables

samanklesaria committed 65 days ago

1d2a672b

Merge pull request #5079 from samanklesaria:eager_sharding_context

a-googler committed 66 days ago

7710c304

Merge pull request #5074 from google:params_metadata_args

a-googler committed 66 days ago

b08cb206

replaced carry_graphdefs -> graphdefs

vfdev-5 committed 66 days ago

aa02f704

Implement *_metadata to other layers

vfdev-5 committed 66 days ago

f320a5c2

Added kernel_metadata/bias_metadata args to nnx layers

vfdev-5 committed 66 days ago

c6036553

Merge pull request #4990 from samanklesaria:issues/4593

a-googler committed 67 days ago

55a23669

Merge pull request #5066 from google:dataclass

a-googler committed 67 days ago

74b72bd5

add nnx dataclass

cgarciae committed 67 days ago

5b4c600a

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

a-googler committed 67 days ago

49ef97c3

Merge pull request #5083 from dan-zheng:docs

a-googler committed 67 days ago

b1f573b7

Update gspmd guide to mention Variable init with eager_sharding

samanklesaria committed 67 days ago

719f6879

Merge branch 'main' into eager_sharding_context

samanklesaria committed 67 days ago

0f5cdf14

Add nnx.use_eager_sharding context manager

samanklesaria committed 67 days ago

f600dcb5

Fixed nnx.scan with carry as pytree and sow

vfdev-5 committed 68 days ago

a611581b

Older