google/flax

Pull Requests Commits

efficient attention implementation with missing files

levskaya committed 5 years ago

0df0f145

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

jheek committed 5 years ago

ee20ca22

use staging jaxpr trace in lazy init

jheek committed 5 years ago

6ccd6a60

fix omnistaging issues

jheek committed 5 years ago

1f553886

Align layernorm dtype handling with batchnorm (i.e., use requested dtype for layernorm outputs, even though intermediate computations are f32).

jekbradbury committed 5 years ago

ed42d067

Updating OSS workflow for docs

a-googler committed 5 years ago

1c509ed4

Merge pull request #328 from lucasb-eyer/patch-3

MostafaDehghani committed 5 years ago

Verified 4578df43

Also add link to Jax's dot_general doc.

lucasb-eyer committed 5 years ago

Verified 23b3be96

Add link to DotGeneral in DenseGeneral doc.

lucasb-eyer committed 5 years ago

Verified 78945d4c

Flax README typos cleanup.

dbieber committed 5 years ago

fbdf8513

Merge pull request #301 from ayush-1506:batchnorm_error_msg

a-googler committed 5 years ago

484cc116

Merge pull request #315 from google:mohitreddy1996-patch-1

a-googler committed 5 years ago

639821e1

Concat the distributed batchnorm to make it more efficient.

a-googler committed 5 years ago

a3a6f7ba

Older