google/flax

Pull Requests Commits

Add __eq__ for variables

samanklesaria committed 52 days ago

1d2a672b

Merge pull request #5079 from samanklesaria:eager_sharding_context

a-googler committed 53 days ago

7710c304

Merge pull request #5074 from google:params_metadata_args

a-googler committed 53 days ago

b08cb206

replaced carry_graphdefs -> graphdefs

vfdev-5 committed 53 days ago

aa02f704

Implement *_metadata to other layers

vfdev-5 committed 53 days ago

f320a5c2

Added kernel_metadata/bias_metadata args to nnx layers

vfdev-5 committed 53 days ago

c6036553

Merge pull request #4990 from samanklesaria:issues/4593

a-googler committed 54 days ago

55a23669

Merge pull request #5066 from google:dataclass

a-googler committed 54 days ago

74b72bd5

add nnx dataclass

cgarciae committed 54 days ago

5b4c600a

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

a-googler committed 54 days ago

49ef97c3

Merge pull request #5083 from dan-zheng:docs

a-googler committed 54 days ago

b1f573b7

Update gspmd guide to mention Variable init with eager_sharding

samanklesaria committed 54 days ago

719f6879

Merge branch 'main' into eager_sharding_context

samanklesaria committed 54 days ago

0f5cdf14

Add nnx.use_eager_sharding context manager

samanklesaria committed 54 days ago

f600dcb5

Fixed nnx.scan with carry as pytree and sow

vfdev-5 committed 55 days ago

a611581b

add promote_dtype to all standard layers

cgarciae committed 55 days ago

3a630ff0

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

a-googler committed 55 days ago

cd37bc90

Merge pull request #5076 from google:update-TestLayersSameGraph

a-googler committed 55 days ago

c19d4c5f

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

a-googler committed 55 days ago

de47d9b7

fix pre-commit issue

cgarciae committed 57 days ago

3daf9532

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

dan-zheng committed 58 days ago

1f67c8b6

Copybara import of the project:

a-googler committed 58 days ago

8e9f8ae3

Copybara import of the project:

cgarciae committed 58 days 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 58 days ago

045df4a3

Added missing LayerNorm test case into TestLayersSameGraph

vfdev-5 committed 58 days ago

dc479ee6

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

a-googler committed 59 days ago

d3754f63

Adding set_mode_info

chapman20j committed 61 days ago

72b8f5d2

support explicit sharding in eager sharding

cgarciae committed 61 days ago

c05d004d

Run pre-commit hook

samanklesaria committed 61 days ago

10c2da40

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

vfdev-5 committed 61 days ago

dcf120b6

Newer Older