google/flax

Pull Requests Commits

Add back "nn" for backward compatibility

marcvanzee committed 4 years ago

3aecf66a

Merge pull request #1630 from jheek:improve-dataclass-doc

a-googler committed 4 years ago

d23a5779

Internal change

a-googler committed 4 years ago

e79a1009

Moves flax.nn --> flax.deprecated.nn and updates all references.

a-googler committed 4 years ago

a79dc33c

Moves flax.nn --> flax.deprecated.nn and updates all references.

marcvanzee committed 4 years ago

96103491

Update flax/struct.py

jheek committed 4 years ago

Verified 9228f8a1

Updates examples/imagenet README.

andsteing committed 4 years ago

8dfcd633

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

Add "smart constructor" example to datataclass docstring

jheek committed 4 years ago

2546724c

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

Older