google/flax

Pull Requests Commits

Merge pull request #5190 from google:hijax-coverage

a-googler committed 12 hours ago

2ed95f31

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 1 day 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 2 days ago

7551bfc3

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

a-googler committed 4 days ago

21695fdf

add HijaxTransformCoverageTest

cgarciae committed 5 days ago

a7397321

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

a-googler committed 6 days ago

3da2276b

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

a-googler committed 6 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 7 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 7 days ago

0ef9abb2

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

a-googler committed 7 days ago

375c63e3

Refactor tabulate

samanklesaria committed 7 days ago

4e7147e9

Use eval_shape for initial call to body

samanklesaria committed 7 days ago

573cd25d

chapman20j committed 8 days ago

772d9974

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

a-googler committed 8 days ago

509780a1

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

a-googler committed 8 days ago

73e7985b

Merge pull request #5183 from samanklesaria/bench_fix

chapman20j committed 8 days ago

Verified 17e2b7e8

Fix bug in graph overhead benchmark

samanklesaria committed 8 days ago

514dba8d

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

ayulockedin committed 8 days ago

903fed7b

Copybara import of the project:

a-googler committed 12 days ago

a7929b37

Merge pull request #5089 from samanklesaria:sharding_metadata

a-googler committed 12 days ago

2863e252

docs: fix typo in description of StatefulLinear

Moriyuki-S committed 12 days ago

b042a8a5

Merge pull request #5133 from samanklesaria:issues/5130

a-googler committed 13 days ago

c4e6bca0

Merge pull request #5166 from ayulockedin:fix-docs-typos

a-googler committed 13 days ago

f42707ab

Merge pull request #5115 from google:hijax-guide

a-googler committed 14 days ago

e8fe693f

improve hijax guide

cgarciae committed 14 days ago

5e31fe9e

Support multiple None and UNCONSTRAINED when resolving logical rules

marksandler2 committed 14 days ago

25610a44

[pmap] Avoid degraded performance under the new `jax.pmap`.

danielsuo committed 15 days ago

5d908519

docs: fix typo 'paramater' -> 'parameter'

ayulockedin committed 15 days ago

15b2f0d7

Merge pull request #5159 from samanklesaria/pr/5140

chapman20j committed 15 days ago

Verified 4cc1d1bb

No public description

oprypin committed 15 days ago

dddd218e

Older