google/flax

Pull Requests Commits

Add function to load checkpoint from absolute path.

anuragarnab committed 3 years ago

917b6bc0

update docs for @nn.compact

a-googler committed 3 years ago

8ff7d702

Merge pull request #2071 from andsteing:mnist

a-googler committed 3 years ago

995bc01f

Makes GPU info shorter (-L flag).

andsteing committed 3 years ago

5a7a345f

Makes annotated_mnist use Optax's xent loss.

andsteing committed 3 years ago

e298d6c1

Remove unused comments related to Python 2 compatibility.

yilei committed 3 years ago

02f53af9

Merge pull request #2064 from marcvanzee:docstring-pooling

a-googler committed 3 years ago

7afd446d

Merge pull request #2039 from marcvanzee:pytype-filter

a-googler committed 3 years ago

22568597

Fixs some docstrings

marcvanzee committed 3 years ago

d4ff240f

Adds default value to pooling.pool)

marcvanzee committed 3 years ago

96f5b6ef

Add support for JAX dynamic stack-based named_call.

levskaya committed 3 years ago

18be4d4d

Bug fixes and generalizations of nn.partitioning api.

levskaya committed 3 years ago

942d6dc2

Removes default_key_shape() from scope.py since jax.random.default_prng_impl is part of the latest JAX release now.

marcvanzee committed 3 years ago

2ac765a5

Merge pull request #2055 from levskaya:dofix

a-googler committed 3 years ago

7371e8b0

Handle edge-case of rate==1.0 in Dropout layer.

levskaya committed 3 years ago

24bc7bd2

Fix linter errors in core/scope.py and core/tracers.py.

marcvanzee committed 3 years ago

c54e1131

Updates `clu.metrics.CollectingMetric` to always use `np.array()`.

andsteing committed 3 years ago

2387439a

Merge pull request #1964 from jheek:module-lifecycle-note

a-googler committed 3 years ago

193f44e9

jheek committed 3 years ago

09bec09f

Module lifecycle note

jheek committed 3 years ago

0238911a

Merge pull request #2036 from marcvanzee:logo

a-googler committed 3 years ago

2c507195

Removes deprecated filterwarning from pytest.ini

marcvanzee committed 3 years ago

68ca372a

Bump Flax/Jax versions in Flax examples.

Dmitry Tsarkov committed 3 years ago

9ce2e302

Adds Flax logo to README.md

marcvanzee committed 3 years ago

a9954a6c

Merge pull request #2038 from google:marcvanzee-patch-13

a-googler committed 3 years ago

4092e467

Update pytest.ini

marcvanzee committed 3 years ago

Verified 181bcaa2

Update pytest.ini

marcvanzee committed 3 years ago

Verified 49b1a44a

Add filterwarning for jax.tree_multimap

marcvanzee committed 3 years ago

Verified d508e94e

Better error messages for loading checkpoints.

lespeholt committed 3 years ago

871358ad

Copybara import of the project:

levskaya committed 3 years ago

3d1b955e

Older