google/flax

Pull Requests Commits

Merge branch 'main' into eager_sharding_context

samanklesaria committed 100 days ago

0f5cdf14

Add nnx.use_eager_sharding context manager

samanklesaria committed 100 days ago

f600dcb5

Fixed nnx.scan with carry as pytree and sow

vfdev-5 committed 100 days ago

a611581b

add promote_dtype to all standard layers

cgarciae committed 100 days ago

3a630ff0

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

a-googler committed 100 days ago

cd37bc90

Merge pull request #5076 from google:update-TestLayersSameGraph

a-googler committed 100 days ago

c19d4c5f

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

a-googler committed 100 days ago

de47d9b7

fix pre-commit issue

cgarciae committed 102 days ago

3daf9532

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

dan-zheng committed 103 days ago

1f67c8b6

Copybara import of the project:

a-googler committed 103 days ago

8e9f8ae3

Copybara import of the project:

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

045df4a3

Added missing LayerNorm test case into TestLayersSameGraph

vfdev-5 committed 104 days ago

dc479ee6

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

a-googler committed 105 days ago

d3754f63

Adding set_mode_info

chapman20j committed 106 days ago

72b8f5d2

support explicit sharding in eager sharding

cgarciae committed 106 days ago

c05d004d

Run pre-commit hook

samanklesaria committed 107 days ago

10c2da40

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

vfdev-5 committed 107 days ago

dcf120b6

Fix trailing whitespace in transforms_test.py

mohsinm-dev committed 107 days ago

b9ac1bce

Remove redundant static args sharding tests

mohsinm-dev committed 107 days ago

6840519c

Fix NNX jit static args with in_shardings issue #4989

mohsinm-dev committed 107 days ago

963f9ba9

Merge pull request #4975 from samanklesaria:issues/4971

a-googler committed 111 days ago

74985b29

Merge pull request #5061 from Lucas-Fernandes-Martins:main

a-googler committed 111 days ago

fe5dec24

Merge pull request #5056 from chapman20j:just_set_mode

a-googler committed 111 days ago

35311878

reverting changes on mnist_tutorial.md

Lucas-Fernandes-Martins committed 112 days ago

21f500ca

fixing documentation on README

Lucas-Fernandes-Martins committed 112 days ago

f2733f13

Remove unused bound method handling code

mohsinm-dev committed 112 days ago

Verified e5e87e4a

[JAX] Replace reference to jax._src.typing with public APIs.

hawkinsp committed 112 days ago

6418755c

fixing Optimizer docstring

Lucas-Fernandes-Martins committed 112 days ago

c3cc9f15

fixing Optimizer docstring

Lucas-Fernandes-Martins committed 112 days ago

2482179d

Newer Older