google/flax

Pull Requests Commits

set axis='batch' for SyncBatchNorm

a-googler committed 2 years ago

0e5ad128

Roll back the Orbax flag flip.

IvyZX committed 2 years ago

751f26a1

Migrate Flax away from `register_keypaths`

IvyZX committed 2 years ago

d129b664

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

a-googler committed 2 years ago

dd14e011

Fix checkpointing docs for Orbax 0.1.4

IvyZX committed 2 years ago

a6b86fb4

Replaces references to jax.numpy.DeviceArray with jax.Array.

hawkinsp committed 2 years ago

dacf0036

Adds metrics function that describes main metrics of the experiment.

andsteing committed 2 years ago

4efd541d

Removes wmt/configs/gpu_mixed_precision.py

andsteing committed 2 years ago

fec41105

Merge pull request #2947 from chiamp:scopeparamshapeerror

a-googler committed 2 years ago

94d50f67

Remove process_index() check in flax/examples, to match Orbax API.

IvyZX committed 2 years ago

594b4529

Filter out orbax temporary checkpoints from the list of all checkpoints.

saxenasaurabh committed 2 years ago

6b4c1075

Eliminate orbax.checkpoint's dependency on Flax by relying on new jax.tree_util functions.

cpgaffney1 committed 2 years ago

f6d81d46

Merge pull request #2939 from google:improve_with_logical_constraint

a-googler committed 2 years ago

351248c8

fix mypy errors

cgarciae committed 2 years ago

5b6579d2

Merge pull request #2931 from google:fix-tabulate

a-googler committed 2 years ago

ccf48a62

ignore FutureWarning from jax

cgarciae committed 2 years ago

d146765a

add rules and mesh docstrings

cgarciae committed 2 years ago

d053913e

permit passing mesh or rules to spmd primitives

cgarciae committed 2 years ago

97668a91

Merge pull request #2946 from IvyZX:push

a-googler committed 2 years ago

d1e420a2

Merge pull request #2934 from 8bitmp3:google-flax-examples

a-googler committed 2 years ago

bdb3f005

Add Flax Google Research examples

8bitmp3 committed 2 years ago

370b1d23

fixed ScopeParamShapeError message

chiamp committed 2 years ago

99ffea08

bump up head version number

IvyZX committed 2 years ago

1418373a

Merge pull request #2944 from IvyZX:push

a-googler committed 2 years ago

323ebc64

IvyZX committed 2 years ago

1b11334a

Allow the checkpoint directory to be nonexistent before saving.

IvyZX committed 2 years ago

e7b83003

Set Orbax as default checkpointing method of Flax. Also added a `DeprecationWarning` (not affecting the tests)

IvyZX committed 2 years ago

8d95e22b

[JAX] Fix up uses of PyTree and PyTreeDef types.

hawkinsp committed 2 years ago

1d669bc5

Migrate Flax dataclass to the newest JAX pytree keypath API.

IvyZX committed 2 years ago

c7fa10a0

cgarciae committed 2 years ago

4c6f0825

Older