google/flax

Pull Requests Commits

allow nnx standalone import

cgarciae committed 73 days ago

6e05d27b

_apply_sharding disallow mixed Explicit/Auto mesh

Cristian Garcia committed 73 days ago

707a569c

Merge pull request #5171 from samanklesaria:issues/5170

a-googler committed 73 days ago

76028b52

allow nnx standalone import p1

Cristian Garcia committed 73 days ago

d1e9f2ef

Merge pull request #5190 from google:hijax-coverage

a-googler committed 75 days ago

2ed95f31

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 76 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 77 days ago

7551bfc3

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

a-googler committed 78 days ago

21695fdf

add HijaxTransformCoverageTest

cgarciae committed 80 days ago

a7397321

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

a-googler committed 81 days ago

3da2276b

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

a-googler committed 81 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 82 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 82 days ago

0ef9abb2

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

a-googler committed 82 days ago

375c63e3

Refactor tabulate

samanklesaria committed 82 days ago

4e7147e9

Use eval_shape for initial call to body

samanklesaria committed 82 days ago

573cd25d

chapman20j committed 83 days ago

772d9974

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

a-googler committed 83 days ago

509780a1

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

a-googler committed 83 days ago

73e7985b

Merge pull request #5183 from samanklesaria/bench_fix

chapman20j committed 83 days ago

Verified 17e2b7e8

Fix bug in graph overhead benchmark

samanklesaria committed 83 days ago

514dba8d

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

ayulockedin committed 83 days ago

903fed7b

Copybara import of the project:

a-googler committed 86 days ago

a7929b37

Merge pull request #5089 from samanklesaria:sharding_metadata

a-googler committed 87 days ago

2863e252

docs: fix typo in description of StatefulLinear

Moriyuki-S committed 87 days ago

b042a8a5

Merge pull request #5133 from samanklesaria:issues/5130

a-googler committed 88 days ago

c4e6bca0

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

a-googler committed 88 days ago

f42707ab

Merge pull request #5115 from google:hijax-guide

a-googler committed 89 days ago

e8fe693f

improve hijax guide

cgarciae committed 89 days ago

5e31fe9e

Support multiple None and UNCONSTRAINED when resolving logical rules

marksandler2 committed 89 days ago

25610a44

Older