google/flax

Pull Requests Commits

Makes flax Modules more compatible with IPython auto-reload.

a-googler committed 356 days ago

e2134afa

Merge pull request #4199 from 8bitmp3:upgrade-nnx-filters

a-googler committed 356 days ago

dd70960a

Merge pull request #4472 from google:nnx-fix-fori

a-googler committed 356 days ago

595e711f

[nnx] test fori_loop output

cgarciae committed 357 days ago

f8164dd4

Merge pull request #4411 from 8bitmp3:update-nnx-merge-docs

a-googler committed 1 year ago

53bde74a

Merge pull request #4325 from 8bitmp3:update-nnx-gemma

a-googler committed 1 year ago

cc29a7b3

Update NNX merge docs in graph.py

8bitmp3 committed 1 year ago

0a40dcdb

Upgrade Flax NNX Gemma

8bitmp3 committed 1 year ago

a46578ec

Introduce a mode for a faster nn.scan mode that avoids extra jax retracing.

levskaya committed 1 year ago

fc38f21e

Merge pull request #4286 from 8bitmp3:update-nnx-jax-transforms

a-googler committed 1 year ago

207966e5

Update Flax NNX vs JAX Transformations guide

8bitmp3 committed 1 year ago

c8bfd2a6

Upgrade NNX Filters guides

8bitmp3 committed 1 year ago

9dd1e4eb

Create sharding via Partitioned.get_sharding()

hhb committed 1 year ago

fcc2e0e5

Merge pull request #4401 from 8bitmp3:update-performance-doc

a-googler committed 1 year ago

554b690b

Merge pull request #4428 from IvyZX:bench

a-googler committed 1 year ago

a785bff2

Merge pull request #4421 from google:nnx-fix-transforms-guide

a-googler committed 1 year ago

f09d1054

Merge pull request #4425 from rajasekharporeddy:testbranch1

a-googler committed 1 year ago

47122f20

add benchmark on state traversal, and a readme

IvyZX committed 1 year ago

f043c5e1

Merge pull request #4364 from Matt-Hurd:matthurd/update_orbax_links

a-googler committed 1 year ago

21585adb

Update links in Why Flax NNX document

rajasekharporeddy committed 1 year ago

c0c4391b

[nnx] fix transforms guide

cgarciae committed 1 year ago

10f34608

Update references to JAX's GitHub repo

jakeharmon8 committed 1 year ago

3e789b71

Merge pull request #4385 from emmanuel-ferdman:main

a-googler committed 1 year ago

30b438c5

Merge pull request #4410 from google:nnx-optimizations-p3

a-googler committed 1 year ago

7eac0514

[nnx] optimize graph

cgarciae committed 1 year ago

217f5ba5

Merge pull request #4415 from tilakrayal:patch-3

a-googler committed 1 year ago

6c840976

Merge pull request #4416 from rajasekharporeddy:main

a-googler committed 1 year ago

27b18291

Remove the repeated methods in flax.nnx.Module documentation

rajasekharporeddy committed 1 year ago

63a1b7d7

Update linen_to_nnx.rst

tilakrayal committed 1 year ago

Verified 4756b34e

Merge pull request #4279 from 8bitmp3:update-nnx-randomness

a-googler committed 1 year ago

a5e62af6

Newer Older