google/flax

Pull Requests Commits

Merge pull request #3314 from jakevdp:typed-prng

a-googler committed 2 years ago

9aeb909b

Merge pull request #3281 from google:method-accept-str-submodules

a-googler committed 2 years ago

a5cf5dbf

allow method argument to accept submodules

cgarciae committed 2 years ago

37cbb7dc

Internal change

romanngg committed 2 years ago

426c0b42

Internal change

romanngg committed 2 years ago

a1829adc

Merge pull request #3317 from google:fix-main3

a-googler committed 2 years ago

23bd52de

cgarciae committed 2 years ago

4f5c143c

add tf-text to doc requirements

cgarciae committed 2 years ago

ad8a20ff

Update handling of typed PRNG keys

jakevdp committed 2 years ago

9a1496fd

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

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

Newer Older