google/flax

Pull Requests Commits

Save types information to allow for fully self-describing checkpoints. Add type-specific metadata to `LazyValue` so that more information can be gleaned without fully restoring the tree.

cpgaffney1 committed 2 years ago

2d70c866

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

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

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

Merge pull request #2976 from google:dependabot/pip/examples/ogbg_molpcba/tensorflow-2.11.1

a-googler committed 2 years ago

34113784

Replace gda_serialization with array_serialization as that's the new name of the library. This is a simple name replacement with no change in functionality.

yashk2810 committed 2 years ago

1145d5e0

Merge pull request #3085 from chiamp:unreferenced_docs

a-googler committed 2 years ago

1785f126

Merge pull request #3079 from chiamp:remove_advanced

a-googler committed 2 years ago

7377d654

Merge pull request #3089 from google:0.6.10

a-googler committed 2 years ago

d1385b0b

cgarciae committed 2 years ago

836aa837

Add test case for degenerate and deterministic attention layer dropout.

levskaya committed 2 years ago

fda64dc4

removed advanced_topics folder and updated link pointers

chiamp committed 2 years ago

67c96403

Older