google/flax

Pull Requests Commits

Merge pull request #5205 from samanklesaria:embed_out_sharding

a-googler committed 69 days ago

2e1f5adc

Merge pull request #5220 from chapman20j:bonsai_link

a-googler committed 69 days ago

93d00808

Pin tqdm version for doctest failures

samanklesaria committed 70 days ago

5bc09657

chapman20j committed 70 days ago

172bafe5

Merge pull request #5204 from chapman20j:view

a-googler committed 71 days ago

4089582f

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

a-googler committed 71 days ago

be338919

chapman20j committed 74 days ago

4f5a64c6

Copybara import of the project:

Cristian Garcia committed 75 days ago

b4e73d79

Fixed mask sharding if inputs is sharded

vfdev-5 committed 76 days ago

0bff68cf

Copybara import of the project:

Cristian Garcia committed 76 days ago

df2ed160

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

danielsuo committed 77 days ago

01c0bd13

Merge pull request #5072 from google:fix-list

a-googler committed 77 days ago

5adb1744

fix List and improve Pytree

cgarciae committed 77 days ago

af408cf5

update flax to version 0.12.3

Cristian Garcia committed 78 days ago

7c487b1e

Revert test to use Exception rather than ValueError

samanklesaria committed 78 days ago

ce2ae486

Apply suggestions from code review

samanklesaria committed 78 days ago

Verified bec643f7

Add out_sharding argument to Embed layer call

samanklesaria committed 78 days ago

ca19b9bc

feat(nnx): add Grouped Query Attention (GQA) support

ayulockedin committed 82 days ago

4daedfef

_apply_sharding disallow mixed Explicit/Auto mesh

Cristian Garcia committed 82 days ago

707a569c

Merge pull request #5171 from samanklesaria:issues/5170

a-googler committed 82 days ago

76028b52

allow nnx standalone import p1

Cristian Garcia committed 83 days ago

d1e9f2ef

Merge pull request #5190 from google:hijax-coverage

a-googler committed 84 days ago

2ed95f31

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 86 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 87 days ago

7551bfc3

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

a-googler committed 88 days ago

21695fdf

add HijaxTransformCoverageTest

cgarciae committed 89 days ago

a7397321

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

a-googler committed 90 days ago

3da2276b

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

a-googler committed 91 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 91 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 91 days ago

0ef9abb2

Newer Older