google/flax

Pull Requests Commits

Merge pull request #4002 from sanderland:fixpickle

a-googler committed 2 years ago

791df52d

Merge pull request #4025 from IvyZX:unbreak

a-googler committed 2 years ago

9b6748aa

Fix Gemma test to unbreak head

IvyZX committed 2 years ago

1c4e3489

Unpin the Orbax pip version.

IvyZX committed 2 years ago

68d5c816

Internal change to fix current head CI

IvyZX committed 2 years ago

bda7da1e

Port gemma/modules to NNX

sibonaci committed 2 years ago

9ca8db56

Support Python 3.9

sibonaci committed 2 years ago

1edfcfa5

Merge pull request #4011 from google:nnx-path-filter

a-googler committed 2 years ago

4d9081e5

Merge pull request #4007 from google:nnx-fix-grad

a-googler committed 2 years ago

54fe15f2

cgarciae committed 2 years ago

1e0a5d10

Port gemma/layers to NNX

sibonaci committed 2 years ago

f62da76f

Helper function for loading params from a linen module

sibonaci committed 2 years ago

ab8c773a

[nnx] add PathContains Filter

cgarciae committed 2 years ago

5ee8e98b

Sander Land committed 2 years ago

f89718e2

Merge pull request #4006 from IvyZX:notftest

a-googler committed 2 years ago

8a2e6413

remove tf from jax_utils_test.py to fix head pytest

IvyZX committed 2 years ago

44e48f9c

Sander Land committed 2 years ago

255564ed

edit flax/errors.py

Sander Land committed 2 years ago

95bad7a1

Sander Land committed 2 years ago

d9731cc3

Cristian Garcia committed 2 years ago

fbdc3394

Merge pull request #3991 from chiamp:iter

a-googler committed 2 years ago

e6b83adc

added nnx.graph.iter_children

chiamp committed 2 years ago

098a2d52

Merge pull request #3983 from chiamp:optimizer

a-googler committed 2 years ago

56216fa6

added wrt option to nnx.Optimizer.update

chiamp committed 2 years ago

9a40b3fa

Merge pull request #3989 from google:nnx-fix-vmap

a-googler committed 2 years ago

551eaf02

Merge pull request #3986 from chiamp:variables

a-googler committed 2 years ago

7a69dcd1

[nnx] vectorize vmap split counts

cgarciae committed 2 years ago

e4f417ac

updated nnx.variables docstrings

chiamp committed 2 years ago

7ca2ee14

Merge pull request #3975 from chiamp:training

a-googler committed 2 years ago

998792a4

Merge pull request #3980 from chiamp:docstrings

a-googler committed 2 years ago

0e5f920e

Older