google/flax

Pull Requests Commits

remove apply overloads

cgarciae committed 2 years ago

2dbb44a2

chiamp committed 2 years ago

53745996

style: pyink reformat

SauravMaheshkar committed 2 years ago

db7b1762

Merge pull request #3265 from peterdavidfagan:typo-fix-dropout

a-googler committed 2 years ago

a7bb19c9

Restore previous _compute_stats behavior in local case.

levskaya committed 2 years ago

976906f9

Merge pull request #3269 from chiamp:release

a-googler committed 2 years ago

d826006e

added cursor api

chiamp committed 2 years ago

925dd1c2

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

Newer Older