google/flax

Pull Requests Commits

Eliminate orbax.checkpoint's dependency on Flax.

cpgaffney1 committed 2 years ago

d9e54817

Added utility to `flax.training.checkpoints` to parse checkpoint folder and return list of checkpointed steps.

Atish Agarwala committed 2 years ago

ce062be2

Updates hyper parameter sweep config in ogbg_molpcba.

andsteing committed 2 years ago

cbd1119b

Delete GDA from flax since it is deprecated

yashk2810 committed 2 years ago

95442f8e

Merge pull request #2919 from chiamp:dict_utility_fns

a-googler committed 2 years ago

779a1226

Merge pull request #2839 from cgarciae:fix-2754

a-googler committed 2 years ago

be3c846c

Merge pull request #2894 from cgarciae:n2883

a-googler committed 2 years ago

8da75423

add 'all' dependency group

cgarciae committed 2 years ago

4364833a

make utility fns work on FrozenDicts

chiamp committed 2 years ago

6597e4cc

Merge pull request #2885 from IvyZX:orbax-upgrade-guide

a-googler committed 2 years ago

3568769a

Update Flax Migrate checkpointing to Orbax guide

8bitmp3 committed 2 years ago

382e4e8d

Add Orbax migration guide

IvyZX committed 2 years ago

f935e541

Merge pull request #2604 from cgarciae:rnn-v2

a-googler committed 2 years ago

10a2123a

Merge pull request #2842 from cgarciae:fix-1322

a-googler committed 2 years ago

f9dab0a8

Merge pull request #2908 from IvyZX:push

a-googler committed 2 years ago

60ab1ccb

install tfds test deps to unblock CI

IvyZX committed 2 years ago

fb010e9a

IvyZX committed 2 years ago

a8e6f5e4

Merge pull request #2898 from chiamp:frozen_dict_tests

a-googler committed 2 years ago

1519161b

migrate tests to regular dicts

chiamp committed 2 years ago

fea91169

cgarciae committed 2 years ago

66dec511

Merge pull request #2860 from cgarciae:check-editable-mode

a-googler committed 2 years ago

776582c3

improve notebook handling

cgarciae committed 2 years ago

d0a50ece

Merge pull request #2897 from chiamp:frozen_dict_utility_fns

a-googler committed 2 years ago

587f6148

Added dict utility functions

chiamp committed 2 years ago

42d51b4f

Merge pull request #2896 from chiamp:frozen_dict_flag_test

a-googler committed 2 years ago

eef5d8b4

Added test for FrozenDict feature flag

chiamp committed 2 years ago

3421c88f

Explicitly mark the lazy-rng with legacy_config, and all new configs should use flax.config.

IvyZX committed 2 years ago

54a205b6

Merge pull request #2872 from chiamp:frozen_dict_feature_flag

a-googler committed 2 years ago

27c37ed1

added feature flag for FrozenDicts

chiamp committed 2 years ago

48f6bda1

Fix or ignore some pytype errors related to jnp.ndarray == jax.Array.

hawkinsp committed 2 years ago

2fbeb3b5

Older