google/flax

Pull Requests Commits

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

Merge pull request #2987 from IvyZX:push

a-googler committed 2 years ago

e061c692

IvyZX committed 2 years ago

b8c89e8c

Merge pull request #2958 from chiamp:pop_doc_str

a-googler committed 2 years ago

8957389d

Merge pull request #2975 from chiamp:frozen_dict_api_doc

a-googler committed 2 years ago

396edea1

Merge pull request #2926 from google:missing-references

a-googler committed 2 years ago

dd8040e5

Copybara import of the project:

dependabot[bot] committed 2 years ago

2248f44b

Copybara import of the project:

dependabot[bot] committed 2 years ago

02e2ad13

Newer Older