google/flax

Pull Requests Commits

Add a TAP test for Flax/Keras integration tests.

divyashreepathihalli committed 98 days ago

3e4a4db0

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

danielsuo committed 99 days ago

01c0bd13

Merge pull request #5072 from google:fix-list

a-googler committed 99 days ago

5adb1744

fix List and improve Pytree

cgarciae committed 99 days ago

af408cf5

update flax to version 0.12.3

Cristian Garcia committed 100 days ago

7c487b1e

_apply_sharding disallow mixed Explicit/Auto mesh

Cristian Garcia committed 104 days ago

707a569c

Merge pull request #5171 from samanklesaria:issues/5170

a-googler committed 104 days ago

76028b52

allow nnx standalone import p1

Cristian Garcia committed 105 days ago

d1e9f2ef

Merge pull request #5190 from google:hijax-coverage

a-googler committed 107 days ago

2ed95f31

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 108 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 109 days ago

7551bfc3

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

a-googler committed 110 days ago

21695fdf

add HijaxTransformCoverageTest

cgarciae committed 112 days ago

a7397321

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

a-googler committed 113 days ago

3da2276b

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

a-googler committed 113 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 113 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 114 days ago

0ef9abb2

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

a-googler committed 114 days ago

375c63e3

Refactor tabulate

samanklesaria committed 114 days ago

4e7147e9

Use eval_shape for initial call to body

samanklesaria committed 114 days ago

573cd25d

chapman20j committed 114 days ago

772d9974

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

a-googler committed 114 days ago

509780a1

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

a-googler committed 114 days ago

73e7985b

Merge pull request #5183 from samanklesaria/bench_fix

chapman20j committed 115 days ago

Verified 17e2b7e8

Fix bug in graph overhead benchmark

samanklesaria committed 115 days ago

514dba8d

docs(nnx): add missing functional args to Conv and LinearGeneral

ayulockedin committed 115 days ago

903fed7b

Copybara import of the project:

a-googler committed 118 days ago

a7929b37

Merge pull request #5089 from samanklesaria:sharding_metadata

a-googler committed 118 days ago

2863e252

docs: fix typo in description of StatefulLinear

Moriyuki-S committed 119 days ago

b042a8a5

Merge pull request #5133 from samanklesaria:issues/5130

a-googler committed 120 days ago

c4e6bca0

Older