google/flax

Pull Requests Commits

marcvanzee committed 4 years ago

a2af2ddc

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

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

Try pinning docutils to 0.16 to fix docs rendering of bulleted lists

avital committed 4 years ago

c107a31f

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

Older