google/flax

Pull Requests Commits

ameya98 committed 4 years ago

284d99b7

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

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

Merge pull request #1545 from jheek:unified-normalization

a-googler committed 4 years ago

550908c0

Merge pull request #1546 from jheek:fix-no-init-hash

a-googler committed 4 years ago

d098f251

jheek committed 4 years ago

6dadc9e0

[JAX] Replace uses of `isinstance(x, jax.numpy.ndarray)` with `isinstance(x, (numpy.ndarray, jax.numpy.ndarray))` where they were leading to test failures.

hawkinsp committed 4 years ago

bbbbdcc9

Unify normalization layers

jheek committed 4 years ago

80c91a0b

Automaticlaly remove non-init attributes from dataclass hash/eq/repr

jheek committed 4 years ago

56f69104

Merge pull request #1547 from jheek/fix-gym-dep

jheek committed 4 years ago

Verified f4586375

Fix gym build using forced downgrade

jheek committed 4 years ago

5315c462

Merge pull request #1537 from levskaya:traceback_filter

a-googler committed 4 years ago

1dfc64b6

Enable traceback filtering of flax internal frames.

levskaya committed 4 years ago

846f951a

Merge pull request #1483 from manifest:PEP-561

a-googler committed 4 years ago

d1343d34

Merge pull request #1535 from jheek:remove-implicit-rank-promotion

a-googler committed 4 years ago

19c4e273

Older