google/flax

Pull Requests Commits

update jax minver to 0.8.1

cgarciae committed 49 days ago

b424c49e

hint that key is cast to string

thijs-vanweezel committed 50 days ago

a7bb109b

test pathcontains with exact=false

thijs-vanweezel committed 50 days ago

12b2e42f

allow substring matching

thijs-vanweezel committed 50 days ago

3d340558

Add is_causal mask argument and tests

ibbyml committed 50 days ago

9bab1485

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

a-googler committed 50 days ago

bee81ed4

Fixed deprecated .value usage failing CI tests

vfdev-5 committed 50 days ago

00cf44e5

Merge pull request #5084 from samanklesaria:variable_eq

a-googler committed 51 days ago

09537f49

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

a-googler committed 51 days ago

b7e66b74

deprecate .value

cgarciae committed 51 days ago

a2eac321

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

a-googler committed 52 days ago

d6e8e14f

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

a-googler committed 52 days ago

544a0505

Merge pull request #5071 from chapman20j:set_mode_info

a-googler committed 52 days ago

0b1d841c

Merge branch 'main' into pr/5055

samanklesaria committed 52 days ago

dece6b34

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

a-googler committed 54 days ago

5109e2c5

cgarciae committed 54 days ago

d982784d

backfill eager_sharding in set_metadata

Cristian Garcia committed 54 days ago

1ac99b09

add HijaxVariable [rollforward]

Cristian Garcia committed 54 days ago

b5db513f

Add __eq__ for variables

samanklesaria committed 56 days ago

1d2a672b

Merge pull request #5079 from samanklesaria:eager_sharding_context

a-googler committed 56 days ago

7710c304

Merge pull request #5074 from google:params_metadata_args

a-googler committed 56 days ago

b08cb206

replaced carry_graphdefs -> graphdefs

vfdev-5 committed 57 days ago

aa02f704

Implement *_metadata to other layers

vfdev-5 committed 57 days ago

f320a5c2

Added kernel_metadata/bias_metadata args to nnx layers

vfdev-5 committed 57 days ago

c6036553

Merge pull request #4990 from samanklesaria:issues/4593

a-googler committed 57 days ago

55a23669

Merge pull request #5066 from google:dataclass

a-googler committed 57 days ago

74b72bd5

add nnx dataclass

cgarciae committed 57 days ago

5b4c600a

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

a-googler committed 57 days ago

49ef97c3

Merge pull request #5083 from dan-zheng:docs

a-googler committed 57 days ago

b1f573b7

Update gspmd guide to mention Variable init with eager_sharding

samanklesaria committed 58 days ago

719f6879

Newer Older