google/flax

Pull Requests Commits

Add -e flag to pip intall

marcvanzee committed 4 years ago

Verified a461117f

Merge pull request #1629 from jheek:lifted-jacobian-products

a-googler committed 4 years ago

03ea1a74

Merge pull request #1640 from jheek/add-nowrap-rtd

jheek committed 4 years ago

Verified 1e2a6c33

jheek committed 4 years ago

0a1c8905

Fix pytype issue

jheek committed 4 years ago

31d04f8b

Add nowrap to RTD

jheek committed 4 years ago

e62c7ccb

Implement vjp and jvp transforms

jheek committed 4 years ago

6906619f

Fixes intendation for some code snippets in `errors.py`.

andsteing committed 4 years ago

38e6bdae

Switch ogbg_molpcba header image to PNG.

ameya98 committed 4 years ago

7eee5b0b

Merge pull request #1623 from jheek:fix-partial-module-method

a-googler committed 4 years ago

93c18dce

Merge pull request #1622 from 8bitmp3/patch-1

andsteing committed 4 years ago

Verified d2fcec57

Add support for transforming methods that are not user functions (e.g.: partial functions).

jheek committed 4 years ago

27b233d7

Lint Flax basics

8bitmp3 committed 4 years ago

Verified 18454b63

Introduce an new system of logical axis names for use with pjit.

levskaya committed 4 years ago

11f8151d

Merge pull request #1611 from patrickvonplaten:patch-1

a-googler committed 4 years ago

8842c69f

Apply suggestions from code review

patrickvonplaten committed 4 years ago

Verified 8faa5f93

Add link to Transformers examples

patrickvonplaten committed 4 years ago

Verified 97524c17

Simplify the logic in prefetch_to_device

josipd committed 4 years ago

48b34ab8

Add '?' to README image URL to invalidate GitHub caching.

ameya98 committed 4 years ago

15addcda

Fix newline in ogbg_molpcba README causing image not to render.

ameya98 committed 4 years ago

5febb766

Fix RNG splitting in core lift functions for KeyArrays.

levskaya committed 4 years ago

6444f97e

Fix a bug in axes_scan for custom prngs.

levskaya committed 4 years ago

2aaff695

Implement jax_utils.replicate using jax.device_put_replicated

jheek committed 4 years ago

97dde9ab

Support new JAX custom PRNG key system.

levskaya committed 4 years ago

a7d55077

Merge pull request #1599 from google/revert-1597-test_399373571

marcvanzee committed 4 years ago

Verified e63b8bd3

Revert "Two small fixes for Flax examples"

marcvanzee committed 4 years ago

Verified 595b17d8

Merge pull request #1582 from avital/fix-bullet-points

avital committed 4 years ago

Verified eea6045d

Merge pull request #1597 from google/test_399373571

marcvanzee committed 4 years ago

Verified ebaf2810

Implements a modular and easily extensible evaluation framework for both TFLite and TFjs. The evaluation framework has the following features:

marcvanzee committed 4 years ago

18207501

Merge pull request #1595 from dan-zheng:fix-typos

a-googler committed 4 years ago

68caa8a2

Older