google/flax

Pull Requests Commits

Allow pass-through implementation of state dict

jheek committed 2 years ago

6fcf8c93

Rolling back google/flax@ba9e24a7b697e6407465cb4b05a24a4cea152248

superbobry committed 2 years ago

7f2426cd

add method migration to Haiku guide

IvyZX committed 2 years ago

cbed59a6

Merge pull request #3290 from chiamp:pjit_logical_partitioning

a-googler committed 2 years ago

15d556f9

Merge pull request #3292 from chiamp:haiku_guide_fix

a-googler committed 2 years ago

ffb68c37

Merge pull request #3293 from google:fix-kw-only-dataclass

a-googler committed 2 years ago

4879b4c0

check number of positional args on Module __init__

cgarciae committed 2 years ago

836e0e43

updated haiku upgrade guide

chiamp committed 2 years ago

031bf3e9

undo recent typing changes as it may be confusing pyright

levskaya committed 2 years ago

ba9e24a7

Fix nn.Module typing.

a-googler committed 2 years ago

79277748

added logical partitioning to pjit guide

chiamp committed 2 years ago

07569f7a

Merge pull request #3288 from IvyZX:orbax

a-googler committed 2 years ago

9e82cdb0

Ignore a very weird pytype error

IvyZX committed 2 years ago

78a97273

Remove string from checkpointing example and unbreak doctest in head

IvyZX committed 2 years ago

c02da8d7

Merge pull request #3286 from google:fix-cloudpickle

a-googler committed 2 years ago

3a9df8e3

fix stackoverflow when loading pickled module

cgarciae committed 2 years ago

ed3ad438

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

Newer Older