google/flax

Pull Requests Commits

vfdev-5 committed 49 days ago

b7dda724

Added mixed precision: computation and weights dtypes

vfdev-5 committed 52 days ago

4fa56ad5

Use grain in gemma example

vfdev-5 committed 52 days ago

1225cad2

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

a-googler committed 54 days ago

3da2276b

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

a-googler committed 54 days ago

13131938

Merge pull request #5141 from samanklesaria:issues/5113

a-googler committed 55 days ago

b4252f89

Update mnist example to use NNX.

danielsuo committed 55 days ago

0ef9abb2

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

a-googler committed 55 days ago

375c63e3

Refactor tabulate

samanklesaria committed 55 days ago

4e7147e9

Use eval_shape for initial call to body

samanklesaria committed 55 days ago

573cd25d

chapman20j committed 56 days ago

772d9974

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

a-googler committed 56 days ago

509780a1

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

a-googler committed 56 days ago

73e7985b

Merge pull request #5183 from samanklesaria/bench_fix

chapman20j committed 56 days ago

Verified 17e2b7e8

Fix bug in graph overhead benchmark

samanklesaria committed 56 days ago

514dba8d

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

ayulockedin committed 56 days ago

903fed7b

Copybara import of the project:

a-googler committed 60 days ago

a7929b37

Merge pull request #5089 from samanklesaria:sharding_metadata

a-googler committed 60 days ago

2863e252

Merge pull request #5133 from samanklesaria:issues/5130

a-googler committed 61 days ago

c4e6bca0

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

a-googler committed 61 days ago

f42707ab

Merge pull request #5115 from google:hijax-guide

a-googler committed 62 days ago

e8fe693f

improve hijax guide

cgarciae committed 62 days ago

5e31fe9e

Support multiple None and UNCONSTRAINED when resolving logical rules

marksandler2 committed 62 days ago

25610a44

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

danielsuo committed 63 days ago

5d908519

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

ayulockedin committed 63 days ago

15b2f0d7

Merge pull request #5159 from samanklesaria/pr/5140

chapman20j committed 63 days ago

Verified 4cc1d1bb

No public description

oprypin committed 63 days ago

dddd218e

fix kw_only_dataclasses for python 3.14 (part 2)

Cristian Garcia committed 63 days ago

dd15749f

samanklesaria committed 63 days ago

ebd78685

flatten respect nnx.data

cgarciae committed 63 days ago

750527aa

Older