google/flax

Pull Requests Commits

Update axes_scan to use extend.linear_util

a-googler committed 2 years ago

451890e9

Merge pull request #3308 from anuragarnab:patch-3

a-googler committed 2 years ago

bec1c80d

[flax] Add module path to nn.module.

JXRiver committed 2 years ago

e3f579a5

Minor fix to jax_utils.py

anuragarnab committed 2 years ago

Verified 72911749

Merge pull request #2585 from cgarciae:rnn-flip

a-googler committed 2 years ago

82dd366e

cgarciae committed 2 years ago

9b2f086d

update min typing_extensions version

levskaya committed 2 years ago

d72ab2c8

Promote dot_general injections from a function to a module.

a-googler committed 2 years ago

1127475b

Allow pass-through implementation of state dict

jheek committed 2 years ago

6fcf8c93

Rolling back google/flax@ba9e24a7b697e6407465cb4b05a24a4cea152248

superbobry committed 2 years ago

7f2426cd

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

Older