google/flax

Pull Requests Commits

implement perturb noop behaviour

cgarciae committed 3 years ago

539cef53

Added builder functions for zeros and ones initializers

chiamp committed 3 years ago

668925d8

Add metadata helper transform for adding axis metadata.

levskaya committed 3 years ago

01460b1c

Restructure Flax RTD

8bitmp3 committed 3 years ago

00d922eb

Merge pull request #2787 from 8bitmp3:fix-indentation

a-googler committed 3 years ago

52be1fdb

Use jax.config to implement flax.config, to support features like runtime overrides. Old config flags are kept as-is to avoid breaking existed tests.

IvyZX committed 3 years ago

a7e43237

Fix indentation in pad_shard_unpad jax_utils.py

8bitmp3 committed 3 years ago

Verified 910b33c7

Fix indentation in Flax ScopeCollectionNotFound errors.py

8bitmp3 committed 3 years ago

Verified e797a9f8

Fix indentation in flax.linen.module

8bitmp3 committed 3 years ago

Verified ddc1d2aa

Add metrics for async checkpoint (beacon and total write duration).

a-googler committed 3 years ago

aa67fb7e

Merge pull request #2785 from jheek:fix-gru-typo

a-googler committed 3 years ago

982c6a46

jheek committed 3 years ago

eeb46a65

Merge pull request #2743 from marcvanzee:scope-tests

a-googler committed 3 years ago

9aa28ac7

vwxyzjn committed 3 years ago

a621fbe6

Reference CleanRL's Efficient Implementation

vwxyzjn committed 3 years ago

e90d6785

Merge branch 'main' of https://www.github.com/google/flax into scope-tests

marcvanzee committed 3 years ago

0f95fae6

rename ConvLSTM to ConvLSTMCell

cgarciae committed 3 years ago

17fae707

Merge pull request #2766 from cgarciae:fix-linen-rst

a-googler committed 3 years ago

62617dbd

Merge pull request #2674 from cgarciae:unbind

a-googler committed 3 years ago

8cd3bf0a

remove incorrect SPMD section in linen.rst

cgarciae committed 3 years ago

db7cb6f0

Minor update flax.linen.Module.unbind

8bitmp3 committed 3 years ago

Verified 22e625d5

Merge branch 'main' into unbind

8bitmp3 committed 3 years ago

Verified a8d5fca0

Minor update flax.linen.Module.unbind

8bitmp3 committed 3 years ago

Verified a6d7704b

Minor update flax.linen.Module.unbind

8bitmp3 committed 3 years ago

Verified e8a58da8

Merge pull request #2756 from chiamp:tabulate

a-googler committed 3 years ago

f3333ce1

Merge pull request #2757 from chiamp:remove_duplicate_flip

a-googler committed 3 years ago

6510a77f

Merge pull request #2759 from cgarciae:fix-logical-names-type

a-googler committed 3 years ago

2d79fdb5

fix logical names

cgarciae committed 3 years ago

d5f367da

removed duplicate flip section in docs

chiamp committed 3 years ago

2221f625

updated tabulate documentation

chiamp committed 3 years ago

220e3ad4

Newer Older