google/flax

Pull Requests Commits

Test JAX custom PRNG system.

levskaya committed 4 years ago

b9697dcd

Merge pull request #1516 from romanngg:patch-1

a-googler committed 4 years ago

02956d81

Move import numpy to top

romanngg committed 4 years ago

Verified a30794bd

Do not call `jnp` on shapes in AutoEncoder example

romanngg committed 4 years ago

Verified ca72d0a9

Merge pull request #1497 from INF800:initializers/add-all-from-jax

a-googler committed 4 years ago

1a117d0a

Merge pull request #1511 from jheek:fix-embedding-hash

a-googler committed 4 years ago

8346ef03

Merge pull request #1499 from levskaya:rematfix

a-googler committed 4 years ago

be522b51

Plumb through the prevent_cse kwarg for jax remat.

levskaya committed 4 years ago

7ed82a49

Fix embedding field annotation

jheek committed 4 years ago

6b0788f8

Improve error reporting for Modules that contain un-hashable attributes

a-googler committed 4 years ago

68ce7afe

Merge pull request #1500 from andsteing/examples

andsteing committed 4 years ago

Verified 78fb4023

Merge pull request #1502 from juntang-zhuang/update_optim_doc

andsteing committed 4 years ago

Verified 69233b5f

Update flax.optim.rst

juntang-zhuang committed 4 years ago

Verified fcfb6ddc

Updates `flax.optim` pydoc.

andsteing committed 4 years ago

ac4e566d

Adds NLP example references.

andsteing committed 4 years ago

191595cb

Add AdaBelief in flax.optim, which adapts stepsize according to "belief" in gradient, and achieves good generalization, fast convergence and training stability.

a-googler committed 4 years ago

bb1f8073

Added glorot and he initializers

rozeappletree committed 4 years ago

Verified a6089931

Adding pod commands and runs to imagenet example

sshahrokhi committed 4 years ago

e30b7f5f

Merge pull request #1479 from jaehlee:main

a-googler committed 4 years ago

9bc7e2a9

add usage note to BatchNorm docstring.

levskaya committed 4 years ago

b2a315c1

Merge pull request #1476 from andsteing/optax

andsteing committed 4 years ago

Verified cb646c74

change base learning_rate to preserve schedule

jaehlee committed 4 years ago

da9d8ef9

Merge pull request #1477 from andsteing:examples

a-googler committed 4 years ago

2ae4e135

Updates link in MNIST notebook #1482.

andsteing committed 4 years ago

0b82c690

Merge pull request #1476 from andsteing:optax

a-googler committed 4 years ago

4f28f6ff

Merge pull request #1456 from jheek:better-transform-error

a-googler committed 4 years ago

8611538e

Merge pull request #1397 from levskaya:variable_transform_fix

a-googler committed 4 years ago

470aaa63

Handle passing Variables and Module instances as arguments in lifted transforms

levskaya committed 4 years ago

b1ca884a

Fix missing "of"

jheek committed 4 years ago

964a92bc

Merge pull request #1478 from jheek:add-dataclass-autocomplete-support

a-googler committed 4 years ago

9350b445

Older