google/flax

Pull Requests Commits

Fixed dataset path in train tests in examples

vfdev-5 committed 356 days ago

f4bb041c

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

a-googler committed 359 days ago

5109e2c5

cgarciae committed 359 days ago

d982784d

backfill eager_sharding in set_metadata

Cristian Garcia committed 359 days ago

1ac99b09

add HijaxVariable [rollforward]

Cristian Garcia committed 359 days ago

b5db513f

Merge pull request #5079 from samanklesaria:eager_sharding_context

a-googler committed 361 days ago

7710c304

Merge pull request #5074 from google:params_metadata_args

a-googler committed 361 days ago

b08cb206

Implement *_metadata to other layers

vfdev-5 committed 361 days ago

f320a5c2

Added kernel_metadata/bias_metadata args to nnx layers

vfdev-5 committed 361 days ago

c6036553

Merge pull request #4990 from samanklesaria:issues/4593

a-googler committed 362 days ago

55a23669

Merge pull request #5066 from google:dataclass

a-googler committed 362 days ago

74b72bd5

add nnx dataclass

cgarciae committed 362 days ago

5b4c600a

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

a-googler committed 362 days ago

49ef97c3

Merge pull request #5083 from dan-zheng:docs

a-googler committed 362 days ago

b1f573b7

Update gspmd guide to mention Variable init with eager_sharding

samanklesaria committed 362 days ago

719f6879

Merge branch 'main' into eager_sharding_context

samanklesaria committed 362 days ago

0f5cdf14

Add nnx.use_eager_sharding context manager

samanklesaria committed 362 days ago

f600dcb5

add promote_dtype to all standard layers

cgarciae committed 363 days ago

3a630ff0

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

a-googler committed 363 days ago

cd37bc90

Merge pull request #5076 from google:update-TestLayersSameGraph

a-googler committed 363 days ago

c19d4c5f

Merge pull request #5070 from google:variable-explicit-sharding-support

a-googler committed 363 days ago

de47d9b7

fix pre-commit issue

cgarciae committed 1 year ago

3daf9532

docs: Document `allow_duplicates` argument of `nnx.to_arrays`.

dan-zheng committed 1 year ago

1f67c8b6

Copybara import of the project:

a-googler committed 1 year ago

8e9f8ae3

Copybara import of the project:

cgarciae committed 1 year ago

18063bea

The default AxisType returned by `make_mesh` is going to change to `AxisType.Explicit` in a follow up CL. This change makes it so that the current code falls back to the old behavior.

yashk2810 committed 1 year ago

045df4a3

Added missing LayerNorm test case into TestLayersSameGraph

vfdev-5 committed 1 year ago

dc479ee6

Merge pull request #4996 from mohsinm-dev:fix-nnx-jit-static-sharding

a-googler committed 1 year ago

d3754f63

support explicit sharding in eager sharding

cgarciae committed 1 year ago

c05d004d

Removed private jax._src.api_util import and copied them into flax

vfdev-5 committed 1 year ago

dcf120b6

Older