google/flax

Pull Requests Commits

Improve Pytree flatten/unflatten

Cristian Garcia committed 2 days ago

e9cb0038

Merge pull request #5204 from chapman20j:view

a-googler committed 2 days ago

4089582f

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

a-googler committed 3 days ago

be338919

chapman20j committed 5 days ago

4f5a64c6

Copybara import of the project:

Cristian Garcia committed 6 days ago

b4e73d79

Fixed mask sharding if inputs is sharded

vfdev-5 committed 7 days ago

0bff68cf

Copybara import of the project:

Cristian Garcia committed 7 days ago

df2ed160

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

danielsuo committed 8 days ago

01c0bd13

Merge pull request #5072 from google:fix-list

a-googler committed 8 days ago

5adb1744

fix List and improve Pytree

cgarciae committed 8 days ago

af408cf5

update flax to version 0.12.3

Cristian Garcia committed 9 days ago

7c487b1e

_apply_sharding disallow mixed Explicit/Auto mesh

Cristian Garcia committed 13 days ago

707a569c

Merge pull request #5171 from samanklesaria:issues/5170

a-googler committed 13 days ago

76028b52

allow nnx standalone import p1

Cristian Garcia committed 14 days ago

d1e9f2ef

Merge pull request #5190 from google:hijax-coverage

a-googler committed 15 days ago

2ed95f31

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 17 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 18 days ago

7551bfc3

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

a-googler committed 19 days ago

21695fdf

add HijaxTransformCoverageTest

cgarciae committed 21 days ago

a7397321

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

a-googler committed 21 days ago

3da2276b

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

a-googler committed 22 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 22 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 23 days ago

0ef9abb2

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

a-googler committed 23 days ago

375c63e3

Refactor tabulate

samanklesaria committed 23 days ago

4e7147e9

Use eval_shape for initial call to body

samanklesaria committed 23 days ago

573cd25d

chapman20j committed 23 days ago

772d9974

Merge pull request #5174 from ayulockedin:doc/audit-conv-layers

a-googler committed 23 days ago

509780a1

Merge pull request #5145 from google:fix-more-typos

a-googler committed 23 days ago

73e7985b

Merge pull request #5183 from samanklesaria/bench_fix

chapman20j committed 23 days ago

Verified 17e2b7e8

Older