google/flax

Pull Requests Commits

refactor RNNCellBase

cgarciae committed 2 years ago

61a44d0d

delete long-unused dotgetter utility

levskaya committed 2 years ago

4f9f64e8

Merge pull request #3155 from IvyZX:push

a-googler committed 2 years ago

fd567c85

IvyZX committed 2 years ago

6ad44b73

Merge pull request #3016 from IvyZX:jit-guide

a-googler committed 2 years ago

bdb68ef5

Merge pull request #3153 from levskaya:trafo_setup_fix

a-googler committed 2 years ago

875b40fe

Merge pull request #3154 from LenaMartens:haiku

a-googler committed 2 years ago

d197ca1a

Apply suggestions from code review

levskaya committed 2 years ago

Verified 98395e99

Update Haiku migration guide.

LenaMartens committed 2 years ago

038d512e

Fix setup trigger issue under sharing and transforms.

levskaya committed 2 years ago

5a82a055

Merge pull request #3139 from google:haiku-migration-guide

a-googler committed 2 years ago

fcd914d8

add haiku_migration_guide

cgarciae committed 2 years ago

ea474fa0

Update Scale up Flax Modules on multiple devices

8bitmp3 committed 2 years ago

87e7fa75

upgrade jit guide

IvyZX committed 2 years ago

acce3a28

Fix Orbax aggregate criteria and bump back to latest Orbax version

IvyZX committed 2 years ago

00a028a4

Merge pull request #3077 from google:eager-bind-submodules

a-googler committed 2 years ago

7b7b8bb5

Merge pull request #3122 from google:levskaya-prelu-1

a-googler committed 2 years ago

e728685c

Add collection to self.scope.reserve(name, col) so that sow works with the same name in different collections.

a-googler committed 2 years ago

266fd52d

Merge pull request #3129 from google:add-checks-to-sequential

a-googler committed 2 years ago

4f403d08

Merge pull request #3131 from google:fix-build

a-googler committed 2 years ago

9e85005b

cgarciae committed 2 years ago

87b33945

cgarciae committed 2 years ago

b94bf54a

use pyproject.toml as hash on CI

cgarciae committed 2 years ago

a0cd96bc

improve docstrings for Sequential

cgarciae committed 2 years ago

4e1835f7

improve error message for Sequential's constructor

cgarciae committed 2 years ago

10f2aff2

Replace uses of deprecated jax.numpy functions:

Jake VanderPlas committed 2 years ago

706063e1

Rolling back [#3104](https://github.com/google/flax/pull/3104) because internal tests are breaking.

chiamp committed 2 years ago

7fbc211b

Flax will default to returning regular dicts instead of FrozenDicts in the future. Some internal-only tests fail because they expect frozen params when creating TrainState.

chiamp committed 2 years ago

59460624

Improve the error message in MultiHeadDotProductAttention.

a-googler committed 2 years ago

a1a669aa

Update activation.py docstring for PReLU

levskaya committed 2 years ago

Verified a7179831

Older