google/flax

Pull Requests Commits

Implement BoxedParam as flax.core.AxisMetadata.

zhangqiaorjc committed 2 years ago

a42b7d2c

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

a-googler committed 2 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 2 years ago

a7e43237

Fix indentation in pad_shard_unpad jax_utils.py

8bitmp3 committed 2 years ago

Verified 910b33c7

Fix indentation in Flax ScopeCollectionNotFound errors.py

8bitmp3 committed 2 years ago

Verified e797a9f8

Fix indentation in flax.linen.module

8bitmp3 committed 2 years ago

Verified ddc1d2aa

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

a-googler committed 2 years ago

aa67fb7e

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

a-googler committed 2 years ago

982c6a46

jheek committed 2 years ago

eeb46a65

Merge pull request #2743 from marcvanzee:scope-tests

a-googler committed 2 years ago

9aa28ac7

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

marcvanzee committed 2 years ago

0f95fae6

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

Merge pull request #2730 from IvyZX:pjit

a-googler committed 3 years ago

0453d38f

pjit guide final touches

IvyZX committed 3 years ago

4920bd85

Update Flax pjit guide notebook

8bitmp3 committed 3 years ago

Verified dcf5517e

create pjit spmd guide doc

IvyZX committed 3 years ago

023511b6

Merge pull request #2691 from chiamp:dict_serialization_errors

a-googler committed 3 years ago

3acb0e0a

Merge pull request #2734 from cgarciae:simplify-axis-rules

a-googler committed 3 years ago

8278a404

Older