google/flax

Pull Requests Commits

Merge pull request #1546 from jheek:fix-no-init-hash

a-googler committed 4 years ago

d098f251

jheek committed 4 years ago

6dadc9e0

[JAX] Replace uses of `isinstance(x, jax.numpy.ndarray)` with `isinstance(x, (numpy.ndarray, jax.numpy.ndarray))` where they were leading to test failures.

hawkinsp committed 4 years ago

bbbbdcc9

Unify normalization layers

jheek committed 4 years ago

80c91a0b

Automaticlaly remove non-init attributes from dataclass hash/eq/repr

jheek committed 4 years ago

56f69104

Merge pull request #1547 from jheek/fix-gym-dep

jheek committed 4 years ago

Verified f4586375

Fix gym build using forced downgrade

jheek committed 4 years ago

5315c462

Merge pull request #1537 from levskaya:traceback_filter

a-googler committed 4 years ago

1dfc64b6

Enable traceback filtering of flax internal frames.

levskaya committed 4 years ago

846f951a

Merge pull request #1483 from manifest:PEP-561

a-googler committed 4 years ago

d1343d34

Merge pull request #1535 from jheek:remove-implicit-rank-promotion

a-googler committed 4 years ago

19c4e273

Merge pull request #1536 from juliuskunze:pixelcnn-fix

a-googler committed 4 years ago

7e4c9456

Fix PixelCNN example

juliuskunze committed 4 years ago

5ae1dfcd

Fix temp sample implicit rank promo

jheek committed 4 years ago

1fc7bfe9

Fix one hot implicit rank promotion

jheek committed 4 years ago

05fe7096

Fix layernorm rank promotion

jheek committed 4 years ago

6f749e00

Fix implicit rank promotion in sst2

jheek committed 4 years ago

bc3ae72a

Fix seq2seq implicit rank promotion

jheek committed 4 years ago

ebe84238

Remove Implicit rank promotion

jheek committed 4 years ago

a401e824

Merge pull request #1531 from jheek:fix-lifting-rng-splits

a-googler committed 4 years ago

cf638d18

Remove users of jax.api.* symbols, in preparation for removing the deprecated jax.api name.

hawkinsp committed 4 years ago

e0a618a9

Remove unnecessary splitting of rngs in lifting logic.

jheek committed 4 years ago

0d69f1a2

Dynamically check whether named call is enabled.

jheek committed 4 years ago

6c46a167

Fix temperature sampler loop termination condition:

renjie-liu committed 4 years ago

4ccd5fff

GPU Mixed precision training for WMT

jheek committed 4 years ago

777ee543

Add optional keep_every_n_steps argument to checkpointing. When enabled, save_checkpoint will keep checkpoints every n steps.

tutmann committed 4 years ago

f92f936b

Merge pull request #1525 from levskaya:modelparamtraversal

a-googler committed 4 years ago

4d7d4580

Merge pull request #1526 from levskaya/pytestinifix

levskaya committed 4 years ago

Verified 37cc870e

Clean up pytest.ini

levskaya committed 4 years ago

965c4fe6

Separate ModelParamTraversal into traverse_util.

levskaya committed 4 years ago

31a1cf8f

Newer Older