google/flax

Pull Requests Commits

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

a-googler committed 2 years ago

8f55865b

Update Flax/Orbax save and load checkpoints

IvyZX committed 2 years ago

6b07e09f

Merge pull request #3022 from chiamp:pretty_repr

a-googler committed 2 years ago

660fd6c1

Added pretty_repr utility fn

chiamp committed 2 years ago

f55c3bf1

Merge pull request #3020 from Mddct:patch-1

a-googler committed 2 years ago

733423a2

Merge pull request #3012 from google:pyproject.toml-setup

a-googler committed 2 years ago

03e738ac

Neilblaze committed 2 years ago

Verified fb4d42bb

modify git repo master to main

Mddct committed 2 years ago

Verified 207b30fb

Merge pull request #3013 from IvyZX:pjit

a-googler committed 2 years ago

332b7d04

fix pspec on replicated array + add sharding utils

IvyZX committed 2 years ago

efe3d57c

Merge pull request #3006 from aniquetahir:main

a-googler committed 2 years ago

93a7de5e

Merge pull request #3010 from google:update-ci-cache

a-googler committed 2 years ago

9f9d48f2

move setup.py to project.toml

cgarciae committed 2 years ago

b2efcd56

Merge pull request #3011 from google:simplify-pyproject.toml

a-googler committed 2 years ago

f486758d

simplify pyproject.toml mypy section

cgarciae committed 2 years ago

f6b817c1

update cache every day

cgarciae committed 2 years ago

9353d0c9

Merge pull request #3007 from andsteing:imagenet_gpu

a-googler committed 2 years ago

322b305f

Bringing changes in https://github.com/google/flax/pull/2955 internally.

chiamp committed 2 years ago

509c02d4

Follow-up to #2985 (adding `config.prefetch`)

andsteing committed 2 years ago

d49183c2

Change Flax's open-source Orbax dependency to build to the orbax-checkpoint domain package, to avoid dependency bloat from unrelated exporting features.

a-googler committed 2 years ago

16b17036

Updates imagenet.ipynb to use GPU Colab runtime.

andsteing committed 2 years ago

8b5744c5

Update module.py

aniquetahir committed 2 years ago

Verified 0fea9e2a

Adds colab_tpu_jax_version.ipynb

andsteing committed 2 years ago

0dde5517

Adds `config.shuffle_buffer_size` and `config.prefetch`.

andsteing committed 2 years ago

71e4432d

Add `kw_only` option to nn.Module dataclass, and a FLIP doc for it.

IvyZX committed 2 years ago

790b275f

Make `restore_args_from_target` API use sharding by default.

IvyZX committed 2 years ago

515393d8

Remove duplicate line in flax imagenet example config

bloops committed 2 years ago

1a605adf

Remove references to deprecated jax.ShapedArray

Jake VanderPlas committed 2 years ago

57a7f26e

Updates examples link.

andsteing committed 2 years ago

5bfa098d

[Trivial] Uptick the version after new release.

IvyZX committed 2 years ago

e17d1cec

Newer Older