google/flax

Pull Requests Commits

cgarciae committed 2 years ago

10768082

Directly call original method if method interceptor stack is empty.

JXRiver committed 2 years ago

a7673630

[flax] Add QK-normalization to MultiHeadDotProductAttention

prazek committed 2 years ago

34452966

Add a Haiku-like methods interceptor API.

JXRiver committed 2 years ago

c8424b04

Merge pull request #3254 from SauravMaheshkar:saurav/vae

a-googler committed 2 years ago

d05aeabd

Merge pull request #3242 from google:improve-typing

a-googler committed 2 years ago

50301491

Merge pull request #3246 from chiamp:cursor

a-googler committed 2 years ago

3ea63811

Merge pull request #3273 from chiamp:bump_version

a-googler committed 2 years ago

634575f0

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

Older