google/flax

Pull Requests Commits

Removed redundant code

vfdev-5 committed 17 days ago

985141fd

Merge pull request #5222 from samanklesaria:doctest_check

a-googler committed 20 days ago

d87f7da7

Merge pull request #5225 from samanklesaria:commit_to_test_ci

a-googler committed 20 days ago

79992cc9

Improve Pytree flatten/unflatten

Cristian Garcia committed 20 days ago

2ab6208e

Pin optax dependency to avoid error

samanklesaria committed 20 days ago

3c2277bf

Merge pull request #5205 from samanklesaria:embed_out_sharding

a-googler committed 22 days ago

2e1f5adc

Merge pull request #5220 from chapman20j:bonsai_link

a-googler committed 22 days ago

93d00808

Pin tqdm version for doctest failures

samanklesaria committed 22 days ago

5bc09657

chapman20j committed 23 days ago

172bafe5

Merge pull request #5204 from chapman20j:view

a-googler committed 24 days ago

4089582f

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

a-googler committed 24 days ago

be338919

chapman20j committed 27 days ago

4f5a64c6

Copybara import of the project:

Cristian Garcia committed 27 days ago

b4e73d79

Fixed mask sharding if inputs is sharded

vfdev-5 committed 29 days ago

0bff68cf

Copybara import of the project:

Cristian Garcia committed 29 days ago

df2ed160

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

danielsuo committed 29 days ago

01c0bd13

Merge pull request #5072 from google:fix-list

a-googler committed 30 days ago

5adb1744

fix List and improve Pytree

cgarciae committed 30 days ago

af408cf5

update flax to version 0.12.3

Cristian Garcia committed 30 days ago

7c487b1e

Revert test to use Exception rather than ValueError

samanklesaria committed 30 days ago

ce2ae486

Apply suggestions from code review

samanklesaria committed 30 days ago

Verified bec643f7

Add out_sharding argument to Embed layer call

samanklesaria committed 30 days ago

ca19b9bc

_apply_sharding disallow mixed Explicit/Auto mesh

Cristian Garcia committed 35 days ago

707a569c

Merge pull request #5171 from samanklesaria:issues/5170

a-googler committed 35 days ago

76028b52

allow nnx standalone import p1

Cristian Garcia committed 36 days ago

d1e9f2ef

Merge pull request #5190 from google:hijax-coverage

a-googler committed 37 days ago

2ed95f31

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 38 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 39 days ago

7551bfc3

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

a-googler committed 41 days ago

21695fdf

add HijaxTransformCoverageTest

cgarciae committed 42 days ago

a7397321

Older