google/flax

Pull Requests Commits

vfdev-5 committed 159 days ago

Verified 5719dd81

vfdev-5 committed 159 days ago

Verified 1e2f9c1a

vfdev-5 committed 159 days ago

Verified 85c198d1

Merge pull request #5204 from chapman20j:view

a-googler committed 160 days ago

4089582f

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

a-googler committed 160 days ago

be338919

chapman20j committed 163 days ago

4f5a64c6

Copybara import of the project:

Cristian Garcia committed 163 days ago

b4e73d79

Fixed mask sharding if inputs is sharded

vfdev-5 committed 165 days ago

0bff68cf

Copybara import of the project:

Cristian Garcia committed 165 days ago

df2ed160

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

danielsuo committed 166 days ago

01c0bd13

Merge pull request #5072 from google:fix-list

a-googler committed 166 days ago

5adb1744

fix List and improve Pytree

cgarciae committed 166 days ago

af408cf5

update flax to version 0.12.3

Cristian Garcia committed 166 days ago

7c487b1e

_apply_sharding disallow mixed Explicit/Auto mesh

Cristian Garcia committed 171 days ago

707a569c

Merge pull request #5171 from samanklesaria:issues/5170

a-googler committed 171 days ago

76028b52

allow nnx standalone import p1

Cristian Garcia committed 172 days ago

d1e9f2ef

Merge pull request #5190 from google:hijax-coverage

a-googler committed 173 days ago

2ed95f31

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 174 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 176 days ago

7551bfc3

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

a-googler committed 177 days ago

21695fdf

add HijaxTransformCoverageTest

cgarciae committed 178 days ago

a7397321

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

a-googler committed 179 days ago

3da2276b

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

a-googler committed 179 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 180 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 180 days ago

0ef9abb2

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

a-googler committed 180 days ago

375c63e3

Refactor tabulate

samanklesaria committed 180 days ago

4e7147e9

Use eval_shape for initial call to body

samanklesaria committed 180 days ago

573cd25d

chapman20j committed 181 days ago

772d9974

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

a-googler committed 181 days ago

509780a1

Older