google/flax

Pull Requests Commits

[flax:examples:mnist] Update mnist example to use NNX.

danielsuo committed 55 days ago

9addd9db

Merge pull request #4975 from samanklesaria:issues/4971

a-googler committed 56 days ago

74985b29

Merge pull request #5061 from Lucas-Fernandes-Martins:main

a-googler committed 56 days ago

fe5dec24

Merge pull request #5056 from chapman20j:just_set_mode

a-googler committed 56 days ago

35311878

reverting changes on mnist_tutorial.md

Lucas-Fernandes-Martins committed 56 days ago

21f500ca

fixing documentation on README

Lucas-Fernandes-Martins committed 56 days ago

f2733f13

[JAX] Replace reference to jax._src.typing with public APIs.

hawkinsp committed 56 days ago

6418755c

fixing Optimizer docstring

Lucas-Fernandes-Martins committed 56 days ago

c3cc9f15

fixing Optimizer docstring

Lucas-Fernandes-Martins committed 56 days ago

2482179d

set mode with tests

chapman20j committed 57 days ago

fd0b62d5

Merge pull request #4999 from samanklesaria:issues/4997

a-googler committed 59 days ago

0183a95e

Merge pull request #5042 from chapman20j:recursive_map

a-googler committed 62 days ago

66508ee1

Merge pull request #5045 from samanklesaria:issues/5044

a-googler committed 62 days ago

dc7e7088

docstring update

chapman20j committed 62 days ago

25b1a25b

Merge pull request #5038 from google:vfdev-5-patch-1

a-googler committed 62 days ago

eb38aad7

Fix shard_map documentation link in compilation.py

vfdev-5 committed 62 days ago

5da4f9c2

Merge pull request #5043 from samanklesaria:issues/4426

a-googler committed 62 days ago

df0cddb3

Use in-place update

samanklesaria committed 62 days ago

2a68940c

samanklesaria committed 62 days ago

da5aca49

Prevent child ctxtag from overwriting parent

samanklesaria committed 62 days ago

fb391c11

filter grad state in nnx.Optimizer

Cristian Garcia committed 63 days ago

241456f6

Merge branch 'google:main' into recursive_map

chapman20j committed 63 days ago

Verified b868f304

Internal change

IvyZX committed 63 days ago

3d1422c9

Internal change

IvyZX committed 63 days ago

60c46d2a

Change introduction to randomness

samanklesaria committed 64 days ago

00601f0d

recursive_map commit

chapman20j committed 65 days ago

0d9cd5aa

chapman20j committed 65 days ago

f0c1a561

Copybara import of the project:

a-googler committed 65 days ago

5a2da6d2

Merge pull request #5037 from google:expose-restore_int_paths

a-googler committed 65 days ago

57c2cc2d

Merge pull request #5032 from google:fix-py314-kw-only-dataclasses

a-googler committed 65 days ago

4d1cfbb1

Older