google/flax

Pull Requests Commits

Allow use of RngBitGenerator TPU hardware PRNG system.

levskaya committed 4 years ago

7b20ed2d

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

Add default and test configs for ogbg_molpcba example.

ameya98 committed 4 years ago

769453ba

Add definition of GNN model for ogbg_molpcba example.

ameya98 committed 4 years ago

cdd77cad

Add input pipeline for ogbg_molpcba example, with unit tests for working with jraph.GraphsTuple.

ameya98 committed 4 years ago

493353c9

Increment jax requirement to fix incompatibility

y0ast committed 4 years ago

Verified 9906f212

fix transform purity check to avoid tree_map.

levskaya committed 4 years ago

01afb539

Merge pull request #1560 from jheek:add-checkpoint-policy-api

a-googler committed 4 years ago

d7172e9d

Merge pull request #1485 from ameya98/patch-1

andsteing committed 4 years ago

Verified 64b2a7a3

Forward experimental checkpoint policy option in lift.remat

jheek committed 4 years ago

bfbaa35a

Merge pull request #1557 from levskaya:rngs_check_fix

a-googler committed 4 years ago

42980115

levskaya committed 4 years ago

8acf39e4

Fix _is_valid_rngs for frozendict of rngs.

levskaya committed 4 years ago

b94311e6

Resubmit normalization refactor

jheek committed 4 years ago

67465368

Merge pull request #1551 from jheek:version-0.3.5

a-googler committed 4 years ago

cbe5ec19

jheek committed 4 years ago

5cd4b7a0

Copybara import of the project:

a-googler committed 4 years ago

480ee046

Older