google/flax

Pull Requests Commits

Reenabled examples tests

vfdev-5 committed 48 days ago

56e4fa9f

Fixed deprecated .value usage failing CI tests

vfdev-5 committed 49 days ago

76cdc322

Merge pull request #5084 from samanklesaria:variable_eq

a-googler committed 49 days ago

09537f49

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

a-googler committed 49 days ago

b7e66b74

deprecate .value

cgarciae committed 49 days ago

a2eac321

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

a-googler committed 49 days ago

d6e8e14f

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

a-googler committed 49 days ago

544a0505

Merge pull request #5071 from chapman20j:set_mode_info

a-googler committed 49 days ago

0b1d841c

Merge branch 'main' into pr/5055

samanklesaria committed 49 days ago

dece6b34

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

a-googler committed 52 days ago

5109e2c5

cgarciae committed 52 days ago

d982784d

backfill eager_sharding in set_metadata

Cristian Garcia committed 52 days ago

1ac99b09

add HijaxVariable [rollforward]

Cristian Garcia committed 52 days ago

b5db513f

Add __eq__ for variables

samanklesaria committed 53 days ago

1d2a672b

Merge pull request #5079 from samanklesaria:eager_sharding_context

a-googler committed 54 days ago

7710c304

Merge pull request #5074 from google:params_metadata_args

a-googler committed 54 days ago

b08cb206

replaced carry_graphdefs -> graphdefs

vfdev-5 committed 54 days ago

aa02f704

Implement *_metadata to other layers

vfdev-5 committed 55 days ago

f320a5c2

Added kernel_metadata/bias_metadata args to nnx layers

vfdev-5 committed 55 days ago

c6036553

Merge pull request #4990 from samanklesaria:issues/4593

a-googler committed 55 days ago

55a23669

Merge pull request #5066 from google:dataclass

a-googler committed 55 days ago

74b72bd5

add nnx dataclass

cgarciae committed 55 days ago

5b4c600a

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

a-googler committed 55 days ago

49ef97c3

Merge pull request #5083 from dan-zheng:docs

a-googler committed 55 days ago

b1f573b7

Update gspmd guide to mention Variable init with eager_sharding

samanklesaria committed 55 days ago

719f6879

Merge branch 'main' into eager_sharding_context

samanklesaria committed 55 days ago

0f5cdf14

Add nnx.use_eager_sharding context manager

samanklesaria committed 55 days ago

f600dcb5

Fixed nnx.scan with carry as pytree and sow

vfdev-5 committed 56 days ago

a611581b

add promote_dtype to all standard layers

cgarciae committed 56 days ago

3a630ff0

Merge pull request #5081 from google:fix-main-2

a-googler committed 56 days ago

cd37bc90

Older