google/flax

Pull Requests Commits

cleanup scan_partial

cgarciae committed 2 years ago

15fea415

add support for initializing carry variables in scan

cgarciae committed 2 years ago

058da51e

Merge pull request #3207 from levskaya:fixgym

a-googler committed 2 years ago

881c4495

Fix warnings from atari gym.

levskaya committed 2 years ago

37e350e4

[JAX] Remove references to jax.config.jax_array.

hawkinsp committed 2 years ago

ce98f350

Fix norm calculation bug for 0-rank arrays.

levskaya committed 2 years ago

3e0d3aa3

[linen] Minor cleanup to normalization code.

chr1sj0nes committed 2 years ago

3ec88c33

[linen] Add alternative, more numerically stable, variance calculation to `LayerNorm`.

chr1sj0nes committed 2 years ago

81467612

[linen] Minor cleanup to normalization code.

levskaya committed 2 years ago

6739a078

Use `import orbax.checkpoint as ocp` to avoid the verbosity of using 'orbax.checkpoint` every time.

cpgaffney1 committed 2 years ago

3e28eeb0

Use `import orbax.checkpoint` which is a better import pattern.

IvyZX committed 2 years ago

9dc4bf4c

Merge pull request #3188 from chiamp:add_dict_migration_index

a-googler committed 2 years ago

5b56a5e5

added dict migration guide to index

chiamp committed 2 years ago

2788ca06

Merge pull request #3187 from google:0.7.0

a-googler committed 2 years ago

7e9d7087

Merge pull request #3186 from google:fix-rnn-docs

a-googler committed 2 years ago

eb1c5ee2

cgarciae committed 2 years ago

d405953d

cgarciae committed 2 years ago

047e4e82

remove cell_size docs from RNN

cgarciae committed 2 years ago

aac2f050

[linen] Only `pmean` the mean if it is non-zero, in normalization modules.

chr1sj0nes committed 2 years ago

0734d006

[linen] Minor cleanup to normalization code.

chr1sj0nes committed 2 years ago

1e7a8b15

Add RMSNorm to documentation

IvyZX committed 2 years ago

1942f3b7

Add propagation of param_dtype to carry initializer.

a-googler committed 2 years ago

8efa9975

Merge pull request #3172 from chiamp:add_absl

a-googler committed 2 years ago

b05c6738

added absltest to linen_recurrent_test

chiamp committed 2 years ago

6bdb1c36

Merge pull request #3167 from chiamp:fix_reference_link

a-googler committed 2 years ago

517225e6

fixed incorrect reference link

chiamp committed 2 years ago

febceabb

Merge pull request #3170 from chiamp:fix_precommit

a-googler committed 2 years ago

3efcd8a2

fix precommit issues

chiamp committed 2 years ago

2d02cc79

Flax is changing the `RNNCellBase` API:

chiamp committed 2 years ago

b6ab7d69

Set default types in Flax for Orbax restoration and add `restore_with_serialized_types` in preparation for an upcoming change.

cpgaffney1 committed 2 years ago

842a2874

Older