google/flax

Pull Requests Commits

Add JAX pytree key path API to tree_utils.

IvyZX committed 2 years ago

596a6267

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

Merge pull request #2881 from chiamp:rmslayernorm

a-googler committed 2 years ago

5f0ac50c

Added RMS normalization layer

chiamp committed 2 years ago

9cff7809

Merge pull request #2873 from levskaya:relaxed_naming2

a-googler committed 2 years ago

98ec8fe2

Refactor and clarify name relaxation policy changes.

levskaya committed 2 years ago

5d8e93db

Merge pull request #2812 from nouiz:imagenet_benchmark

a-googler committed 2 years ago

018725f3

Merge pull request #2879 from cpgaffney1:dev_branch

a-googler committed 2 years ago

5d4040a4

Add additional Orbax link to use_checkpointing.ipynb

cpgaffney1 committed 2 years ago

Verified 21ed0fbd

Add additional Orbax link to use_checkpointing.md

cpgaffney1 committed 2 years ago

Verified 75db8675

Older