google/flax

Pull Requests Commits

cgarciae committed 2 years ago

68b334e7

removed FrozenDict section

chiamp committed 2 years ago

f4386e41

fix carry slice logic

cgarciae committed 2 years ago

50f492b2

[linen] Use `stack` instead of `concatenate` in `compute_stats`, to handle scalar stats case.

chr1sj0nes committed 2 years ago

15d6857e

Fix typo in struct.py documentation

marcoselvi committed 2 years ago

Verified 0a4b0d68

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

Update docs/guides/haiku_migration_guide.rst

cgarciae committed 2 years ago

Verified 06583b50

Update docs/guides/haiku_migration_guide.rst

cgarciae committed 2 years ago

Verified ff9a9d9f

Update docs/guides/haiku_migration_guide.rst

cgarciae committed 2 years ago

Verified b3aaa7c9

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

chr1sj0nes committed 2 years ago

81467612

adds table_kwargs and column_kwargs arguments to tabulate

cgarciae committed 2 years ago

c5f41520

[linen] Minor cleanup to normalization code.

levskaya committed 2 years ago

6739a078

add scan over layers section

cgarciae committed 2 years ago

18d021f1

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

Newer Older