google/flax

Pull Requests Commits

add multi-scope lifting

jheek committed 5 years ago

96d5a812

Revert flax.nn changes

jheek committed 5 years ago

3d9d183b

functional core

jheek committed 5 years ago

369da542

Changed hardcoded type "float32" to dtype (models.py)

a-googler committed 5 years ago

a5dfa290

Remove expectations on secs_per_epoch metric.

a-googler committed 5 years ago

0f66ea70

Merge pull request #352 from wrzadkow:assert-lamb

a-googler committed 5 years ago

8ac99b8f

Merge pull request #358 from juliakreutzer:master

a-googler committed 5 years ago

b69f4038

Revert "Internal only change"

avital committed 5 years ago

5543fae9

Merge pull request #341 from DanGraur/fix-example-imgnet-training

avital committed 5 years ago

Verified 432d5ee2

extracted hard-coded hyperparameters into flags

juliakreutzer committed 5 years ago

3f33f80e

Start exporting event files to benchmark_output_dir as well, as this will help with easier integration with tensorboard.dev

a-googler committed 5 years ago

c4725328

Internal only change

a-googler committed 5 years ago

d3b6b359

Assert valid learning rate in LAMB optimizer

wrzadkow committed 5 years ago

5f7c2e26

Merge pull request #316 from danielsuo/dsuo/improve-check-howto-error

avital committed 5 years ago

Verified 8346a106

Update check_howto_diffs.sh

avital committed 5 years ago

Verified 4da1a18c

Update check_howto_diffs.sh

avital committed 5 years ago

Verified 09241205

Merge branch 'master' into dsuo/improve-check-howto-error

avital committed 5 years ago

Verified 31f8b37e

Merge pull request #313 from danielsuo/dsuo/howto-individual-pages

avital committed 5 years ago

Verified f4111921

Merge pull request #346 from DanGraur:jax-utils-assert-message-fix

a-googler committed 5 years ago

15223579

Start exporting parameters used by the benchmarks.

a-googler committed 5 years ago

060737d4

removed slashes from improved _prefetch assert.

levskaya committed 5 years ago

Verified 2e7012da

Avoid creating an array copy in the replicate function if possible.

akolesnikoff committed 5 years ago

8945bb83

a-googler committed 5 years ago

e3f3aa3f

Merge pull request #347 from DanGraur:optim-init-fix

a-googler committed 5 years ago

5bbf0ac8

Fixed a bug in the optim package, wherein the ModelParamTraversal and DynamicScale modules were not public.

Dan Graur committed 5 years ago

e7dbb82a

use more efficient softmax in dot product attention.

levskaya committed 5 years ago

7fcc7f6c

Added an error message to the assert in jax_utils.prefetch_to_device(*).

Dan Graur committed 5 years ago

1a75efc0

fix flax autoregressive cache initialization for upcoming JAX changes.

levskaya committed 5 years ago

b566e59f

a-googler committed 5 years ago

a26d5ecc

Merge pull request #335 from jheek:omnistaging-fixes

a-googler committed 5 years ago

daff4030

Older