google/flax

Pull Requests Commits

Raise informative error if user tries to naively pass a Module to a jitted function.

Melissa Tan committed 3 years ago

804a165a

Plumb spmd_axis_name through transforms.vmap through to JAX vmap

James Lee-Thorp committed 3 years ago

07e513f6

Adds axis_name and axis_index_groups to LayerNorm and GroupNorm.

a-googler committed 3 years ago

2b73efdd

Merge pull request #2407 from yklcs:bump-rich

a-googler committed 3 years ago

c3e5363d

Merge pull request #2405 from andsteing:examples

a-googler committed 3 years ago

5c09c3a3

Ease rich version requirement

yklcs committed 3 years ago

72189153

Updates examples/{imagenet,wmt} requirements.

andsteing committed 3 years ago

f175457a

Removes `flax.optim` & bumps version to 0.6.0

andsteing committed 3 years ago

cda7a4c8

Merge pull request #2392 from banda-larga:patch-1

a-googler committed 3 years ago

47418b32

Fix minor typos in struct.py docstring.

a-googler committed 3 years ago

99294fb4

fixed jax_for_the_impatient line

banda-larga committed 3 years ago

43c3c29b

Plumb spmd_axis_name from vmap_with_axes through to JAX vmap

James Lee-Thorp committed 3 years ago

8fe8f5eb

Update minimal JAX version to latest (0.3.16).

James Lee-Thorp committed 3 years ago

d9897c59

banda-larga committed 3 years ago

3d77db1e

Removes `flax.optim.dynamic_scale`.

andsteing committed 3 years ago

517b7635

Update minimal JAX version

James Lee-Thorp committed 3 years ago

28291a89

Merge pull request #2389 from mattjj:skip-buggy-test

a-googler committed 3 years ago

890a3913

skip remat test that fails with autodiff

mattjj committed 3 years ago

9b06e96c

Merge pull request #2372 from cgarciae:test-docstrings

a-googler committed 3 years ago

f7545411

Merge pull request #2374 from levskaya:doctweaks2

a-googler committed 3 years ago

2f4628c9

remove whitespace

cgarciae committed 3 years ago

a04384f6

Add option `auto_flush` to `flax.metrics.tensorboard.SummaryWriter`

duckworthd committed 3 years ago

ba5ce6ce

Copies `dynamic_scale.py` from `optim/` to `training/`.

andsteing committed 3 years ago

868d3595

Allow linen's Conv layer to operate on arbitrary-rank inputs.

a-googler committed 3 years ago

e1712906

Merge pull request #2369 from Jeevesh8:patch-1

a-googler committed 3 years ago

4a0f78ac

updated supported transforms in lifting docs.

levskaya committed 3 years ago

6d561493

test docstrings with autodoc on CI

cgarciae committed 3 years ago

099642e3

Merge pull request #2373 from levskaya:doctweaks

a-googler committed 3 years ago

131fcf62

Add some docstrings to the old flax.training.common_utils module.

levskaya committed 3 years ago

0b960eb8

Update ``state_params.md`` alongside ``state_params.ipynb``

Jeevesh Juneja committed 3 years ago

Verified cae20c38

Older