google/flax

Pull Requests Commits

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

Update VAE example

canyon289 committed 2 years ago

db97574b

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

eagerly bind submodules

cgarciae committed 2 years ago

8f122f60

Merge pull request #3118 from google:imagenet-flax-bump

a-googler committed 2 years ago

df66c812

Merge pull request #3095 from chiamp:fix-missing-seperator

a-googler committed 2 years ago

f85c575c

fixed missing separator for rng fold in

chiamp committed 2 years ago

eb82a6cc

Bump imagenet Flax to 0.6.5

IvyZX committed 2 years ago

Verified 667a344f

Merge pull request #3114 from google:explicit-dropout-rng

a-googler committed 2 years ago

b87f5b32

allow manually specifying the rng key for Dropout

cgarciae committed 2 years ago

4bed5e92

added dict migration guide

chiamp committed 2 years ago

e48a2912

Set orbax-as-backend as default for checkpointing.

IvyZX committed 2 years ago

380fef6d

Copybara import of the project:

levskaya committed 2 years ago

63df54ad

Merge pull request #2901 from google:cgarciae-getattribute-no-recursive-register

a-googler committed 2 years ago

cc0aada0

Newer Older