google/flax

Pull Requests Commits

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

a-googler committed 4 years ago

5012a650

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

[JAX] Replace uses of deprecated `jax.ops.index_update(x, idx, y)` APIs with their up-to-date, more succinct equivalent `x.at[idx].set(y)`.

hawkinsp committed 4 years ago

facbc5d7

dan-zheng committed 4 years ago

86372e3e

[JAX] Update JAX users in preparation for a change that makes iteration over a JAX array return JAX arrays, instead of NumPy arrays.

hawkinsp committed 4 years ago

b391239d

Merge pull request #1581 from jheek:map-variables-tranfsorm

a-googler committed 4 years ago

cba70922

Map variables for linen

jheek committed 4 years ago

45a3f081

Fix links in Colab and README for ogbg_molpcba example.

ameya98 committed 4 years ago

9e9fc29a

Merge pull request #1583 from y0ast/patch-1

jheek committed 4 years ago

Verified fb5d2c59

Add BUILD, OWNERS and requirements.txt for ogbg_molpcba example.

ameya98 committed 4 years ago

9db5a5e4

Add README for ogbg_molpcba example.

ameya98 committed 4 years ago

ff83d0ec

Add Colab notebook for demonstrating ogbg_molpcba training.

ameya98 committed 4 years ago

ca9325b2

Add benchmark for ogbg_molpcba example.

ameya98 committed 4 years ago

472d23ec

Add main.py as the entry point for running the ogbg_molpcba example.

ameya98 committed 4 years ago

76bb33ab

Add training and evaluation code, with unit tests.

ameya98 committed 4 years ago

18d40330

Older