google/flax

Pull Requests Commits

levskaya committed 4 years ago

8acf39e4

Fix _is_valid_rngs for frozendict of rngs.

levskaya committed 4 years ago

b94311e6

Resubmit normalization refactor

jheek committed 4 years ago

67465368

Merge pull request #1551 from jheek:version-0.3.5

a-googler committed 4 years ago

cbe5ec19

jheek committed 4 years ago

5cd4b7a0

Copybara import of the project:

a-googler committed 4 years ago

480ee046

Merge pull request #1545 from jheek:unified-normalization

a-googler committed 4 years ago

550908c0

Merge pull request #1546 from jheek:fix-no-init-hash

a-googler committed 4 years ago

d098f251

jheek committed 4 years ago

6dadc9e0

[JAX] Replace uses of `isinstance(x, jax.numpy.ndarray)` with `isinstance(x, (numpy.ndarray, jax.numpy.ndarray))` where they were leading to test failures.

hawkinsp committed 4 years ago

bbbbdcc9

Unify normalization layers

jheek committed 4 years ago

80c91a0b

Automaticlaly remove non-init attributes from dataclass hash/eq/repr

jheek committed 4 years ago

56f69104

Merge pull request #1547 from jheek/fix-gym-dep

jheek committed 4 years ago

Verified f4586375

Fix gym build using forced downgrade

jheek committed 4 years ago

5315c462

Merge pull request #1537 from levskaya:traceback_filter

a-googler committed 4 years ago

1dfc64b6

Enable traceback filtering of flax internal frames.

levskaya committed 4 years ago

846f951a

Merge pull request #1483 from manifest:PEP-561

a-googler committed 4 years ago

d1343d34

Merge pull request #1535 from jheek:remove-implicit-rank-promotion

a-googler committed 4 years ago

19c4e273

Merge pull request #1536 from juliuskunze:pixelcnn-fix

a-googler committed 4 years ago

7e4c9456

Fix PixelCNN example

juliuskunze committed 4 years ago

5ae1dfcd

Fix temp sample implicit rank promo

jheek committed 4 years ago

1fc7bfe9

Fix one hot implicit rank promotion

jheek committed 4 years ago

05fe7096

Fix layernorm rank promotion

jheek committed 4 years ago

6f749e00

Fix implicit rank promotion in sst2

jheek committed 4 years ago

bc3ae72a

Fix seq2seq implicit rank promotion

jheek committed 4 years ago

ebe84238

Remove Implicit rank promotion

jheek committed 4 years ago

a401e824

Merge pull request #1531 from jheek:fix-lifting-rng-splits

a-googler committed 4 years ago

cf638d18

Remove users of jax.api.* symbols, in preparation for removing the deprecated jax.api name.

hawkinsp committed 4 years ago

e0a618a9

Remove unnecessary splitting of rngs in lifting logic.

jheek committed 4 years ago

0d69f1a2

Dynamically check whether named call is enabled.

jheek committed 4 years ago

6c46a167

Newer Older