google/flax

Pull Requests Commits

chiamp committed 2 years ago

5eeb5066

Merge pull request #3258 from chiamp:release2

a-googler committed 2 years ago

fe54d391

fix typo neutral -> neural

peterdavidfagan committed 2 years ago

f6427e2d

bump version to 0.7.2

chiamp committed 2 years ago

f4bc733b

Merge pull request #3241 from PhilipVinc:patch-3

a-googler committed 2 years ago

d883d34c

[JAX] Replace uses of jax.experimental.pjit.with_sharding_constraint with jax.lax.with_sharding_constraint.

hawkinsp committed 2 years ago

8e51b714

Add use_fast_variance option to GroupNorm and BatchNorm to allow disabling it.

a-googler committed 2 years ago

b964fb9a

feat: add configs for vae example

SauravMaheshkar committed 2 years ago

159a3aa7

Fix nn.Module typing.

levskaya committed 2 years ago

8da8c46c

cgarciae committed 2 years ago

36e51418

add typing overloads for variable

cgarciae committed 2 years ago

8525a9b7

Expose the "name" attribute for type checkers.

ashishenoyp committed 2 years ago

a6ed700c

Use field_specifiers instead of field_descriptors in @dataclass_transform.

a-googler committed 2 years ago

9b0b4efd

Expose the "name" attribute for type checkers.

a-googler committed 2 years ago

258d56fc

PhilipVinc committed 2 years ago

81c274af

Apply suggestions from code review

PhilipVinc committed 2 years ago

Verified a30f09f1

PhilipVinc committed 2 years ago

62efe127

make flax.core.copy `add_or_replace` optional

PhilipVinc committed 2 years ago

940ff5d9

add overload signatures for Module.apply

cgarciae committed 2 years ago

1e2d1c71

add overload signatures for param

cgarciae committed 2 years ago

d69d106d

Merge pull request #3234 from chiamp:release

a-googler committed 2 years ago

675e34da

chiamp committed 2 years ago

88934902

On top of #3217, make sure the shape compatibility works before the next Orbax release.

IvyZX committed 2 years ago

3741278d

Merge pull request #3231 from google:fix-scan-docs

a-googler committed 2 years ago

f1032336

improve scan docs

cgarciae committed 2 years ago

93d48c19

Let Flax-Orbax to not port the shape of `target` arrays when they port the `target` shardings.

IvyZX committed 2 years ago

719217b6

Merge pull request #3223 from IvyZX:orbax

a-googler committed 2 years ago

2497f823

Fix checkpointing guide

IvyZX committed 2 years ago

2857f30f

Merge pull request #3195 from google:haiku-migration-scan-over-layers

a-googler committed 2 years ago

9ba7cc0c

Merge pull request #3219 from google:fix-main

a-googler committed 2 years ago

4f1884c7

Newer Older