google/flax

Pull Requests Commits

Add vmap + BatchNorm example to the state howto

JamesAllingham committed 4 years ago

00e9fea7

Use notify_all instead of notifyAll method that was deprecated in Python 3.10.

tirkarthi committed 4 years ago

c1e5bf62

Merge pull request #1306 from bastings:sst2_optax

a-googler committed 4 years ago

48b5707a

Merge pull request #1291 from avital:attention-weights

a-googler committed 4 years ago

e95acab1

Add more assertions to dot_product_attention_weights

avital committed 4 years ago

710d2023

bastings committed 4 years ago

a64d88d5

Move to use Optax optimizer.

bastings committed 4 years ago

afd14928

Merge remote-tracking branch 'upstream/master' into sst2_optax

bastings committed 4 years ago

cf27baf3

Correct assertion string

avital committed 4 years ago

Verified 43a98165

jheek committed 4 years ago

1a0a0136

Renames jax.{host_id,host_count} -> jax.{process_index,process_count} & bumps jaxlib version.

andsteing committed 4 years ago

d969e64d

Fixes type examples/mnist/README.md

andsteing committed 4 years ago

Verified ecf4c592

Switch to Optax optimizer and use TrainState.

bastings committed 4 years ago

a3990ed1

jheek committed 4 years ago

e62dca6f

Merge pull request #1295 from n2cholas:scan-docs

a-googler committed 4 years ago

0082b90c

Update flax/linen/transforms.py

n2cholas committed 4 years ago

de51cf3d

Fix linen.scan docstring example.

n2cholas committed 4 years ago

dc78886b

Updates MNIST from Colab.

andsteing committed 4 years ago

9ce62581

incorporate feedback

jheek committed 4 years ago

efad2fa1

Merge pull request #1301 from google/andsteing-patch-1

jheek committed 4 years ago

Verified 63bd1339

Updates examples/mnist to use Optax.

andsteing committed 4 years ago

c4b34644

Ignore `jax.host_{id,count}()` deprecation warning

andsteing committed 4 years ago

Verified 2a459324

Merge pull request #1300 from google/andsteing-patch-1

jheek committed 4 years ago

Verified 87a483b2

Bumps JAX version to 0.2.13

andsteing committed 4 years ago

Verified a9e3f3cb

Merge pull request #1280 from google/marcvanzee-patch-10

marcvanzee committed 4 years ago

Verified d804b906

Merge pull request #1274 from jheek/update-install-instructions

marcvanzee committed 4 years ago

Verified ac75341b

Merge pull request #1281 from yashs97/fix-contributing-doc

marcvanzee committed 4 years ago

Verified 4f5438e0

Shows boilerplate cell by default to make sure it's executed.

andsteing committed 4 years ago

c6877e94

[flex/linen]: Make `rngs` type "Optional" to avoid pytype errors.

liviosoares committed 4 years ago

0848d9cf

Merge pull request #1284 from jheek:init-mutable

a-googler committed 4 years ago

2faf0eaf

Newer Older