google/flax

Pull Requests Commits

enable mlp test

Cristian Garcia committed 29 days ago

d9dee2d2

Merge pull request #5205 from samanklesaria:embed_out_sharding

a-googler committed 30 days ago

2e1f5adc

Merge pull request #5220 from chapman20j:bonsai_link

a-googler committed 30 days ago

93d00808

chapman20j committed 31 days ago

172bafe5

Merge pull request #5204 from chapman20j:view

a-googler committed 32 days ago

4089582f

Merge pull request #5212 from google:fix-dropout-5209

a-googler committed 32 days ago

be338919

chapman20j committed 35 days ago

4f5a64c6

Copybara import of the project:

Cristian Garcia committed 35 days ago

b4e73d79

Fixed mask sharding if inputs is sharded

vfdev-5 committed 37 days ago

0bff68cf

Copybara import of the project:

Cristian Garcia committed 37 days ago

df2ed160

[pmap] Speed up the new `jax.pmap` implementation.

danielsuo committed 38 days ago

01c0bd13

Merge pull request #5072 from google:fix-list

a-googler committed 38 days ago

5adb1744

fix List and improve Pytree

cgarciae committed 38 days ago

af408cf5

update flax to version 0.12.3

Cristian Garcia committed 38 days ago

7c487b1e

Revert test to use Exception rather than ValueError

samanklesaria committed 38 days ago

ce2ae486

Apply suggestions from code review

samanklesaria committed 38 days ago

Verified bec643f7

Add out_sharding argument to Embed layer call

samanklesaria committed 38 days ago

ca19b9bc

_apply_sharding disallow mixed Explicit/Auto mesh

Cristian Garcia committed 43 days ago

707a569c

Merge pull request #5171 from samanklesaria:issues/5170

a-googler committed 43 days ago

76028b52

allow nnx standalone import p1

Cristian Garcia committed 44 days ago

d1e9f2ef

Merge pull request #5190 from google:hijax-coverage

a-googler committed 45 days ago

2ed95f31

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 46 days ago

7e5c70a0

Shard split rngs in lifted vmap if spmd_axis_name is given and applied to vmapped axis. Without this, jax.vmap gives an error: `ValueError: Mapped away dimension of inputs passed to vmap should be sharded the same. Got inconsistent axis specs: None vs batch` due to split_rngs being replicated.

jossb-iso committed 47 days ago

7551bfc3

Merge pull request #5157 from Moriyuki-S:docs/fix-nnx-functional-api-typo

a-googler committed 49 days ago

21695fdf

add HijaxTransformCoverageTest

cgarciae committed 50 days ago

a7397321

Merge pull request #5188 from samanklesaria:pr/5002_copy

a-googler committed 51 days ago

3da2276b

Merge pull request #5142 from google:docs-typo-fix-_cached_partial

a-googler committed 51 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 52 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 52 days ago

0ef9abb2

Merge pull request #5186 from samanklesaria:issues/5067_copy

a-googler committed 52 days ago

375c63e3

Older