google/flax

Pull Requests Commits

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

Removed a superfluous assignment in the examples/imagenet/train.py source.

Dan Graur committed 5 years ago

4c2319fd

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

Export model_name from the MNIST benchmark.

a-googler committed 5 years ago

59e480f7

Export model_name from the MNIST benchmark

a-googler committed 5 years ago

81752d83

Remove loss scaling from README

mohitreddy1996 committed 5 years ago

Verified 95844ce2

Make clear that we are undoing git apply if tests fail in check_howto_diffs.sh

danielsuo committed 5 years ago

b8610d70

Improve error message in check_howto_diffs.sh

danielsuo committed 5 years ago

7c649dd1

Remove loss_scaling flag from imagenet benchmark

mohitreddy1996 committed 5 years ago

Verified 5a1787b3

Move howto index page

danielsuo committed 5 years ago

f62618c1

Add new howto docs

danielsuo committed 5 years ago

7d5fdc36

Newer Older