google/flax

Pull Requests Commits

vfdev-5 committed 42 days ago

b7dda724

Added mixed precision: computation and weights dtypes

vfdev-5 committed 45 days ago

4fa56ad5

Use grain in gemma example

vfdev-5 committed 45 days ago

1225cad2

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

a-googler committed 47 days ago

3da2276b

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

a-googler committed 47 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 48 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 48 days ago

0ef9abb2

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

a-googler committed 48 days ago

375c63e3

Refactor tabulate

samanklesaria committed 48 days ago

4e7147e9

Use eval_shape for initial call to body

samanklesaria committed 48 days ago

573cd25d

chapman20j committed 49 days ago

772d9974

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

a-googler committed 49 days ago

509780a1

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

a-googler committed 49 days ago

73e7985b

Merge pull request #5183 from samanklesaria/bench_fix

chapman20j committed 49 days ago

Verified 17e2b7e8

Fix bug in graph overhead benchmark

samanklesaria committed 49 days ago

514dba8d

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

ayulockedin committed 49 days ago

903fed7b

Copybara import of the project:

a-googler committed 53 days ago

a7929b37

Merge pull request #5089 from samanklesaria:sharding_metadata

a-googler committed 53 days ago

2863e252

Merge pull request #5133 from samanklesaria:issues/5130

a-googler committed 54 days ago

c4e6bca0

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

a-googler committed 54 days ago

f42707ab

Merge pull request #5115 from google:hijax-guide

a-googler committed 55 days ago

e8fe693f

improve hijax guide

cgarciae committed 55 days ago

5e31fe9e

Support multiple None and UNCONSTRAINED when resolving logical rules

marksandler2 committed 55 days ago

25610a44

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

danielsuo committed 56 days ago

5d908519

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

ayulockedin committed 56 days ago

15b2f0d7

Merge pull request #5159 from samanklesaria/pr/5140

chapman20j committed 56 days ago

Verified 4cc1d1bb

No public description

oprypin committed 56 days ago

dddd218e

fix kw_only_dataclasses for python 3.14 (part 2)

Cristian Garcia committed 56 days ago

dd15749f

samanklesaria committed 56 days ago

ebd78685

flatten respect nnx.data

cgarciae committed 56 days ago

750527aa

Older