google/flax

Pull Requests Commits

cgarciae committed 356 days ago

1c223946

[nnx] add flaxlib

cgarciae committed 356 days ago

59d97246

[nnx] cache flatten

cgarciae committed 359 days ago

cf6db718

[nnx] simpllify unflatten

cgarciae committed 1 year ago

e9c635d7

cgarciae committed 1 year ago

6cf5b7d3

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

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

Older