google/flax

Pull Requests Commits

improve splitting static and dynamic arguments

cgarciae committed 2 years ago

f90e5756

Report forward and backward pass FLOPs of modules and submodules in `linen.Module.tabulate` and `summary.tabulate` (in new `flops` and `vjp_flops` table columns)

romanngg committed 2 years ago

bf2faf57

Merge pull request #3343 from chiamp:haiku_upgrade

a-googler committed 2 years ago

5d846a59

updated haiku guide with new JAX RNG api

chiamp committed 2 years ago

089f237b

Migrate flax from using old-style PRNG keys to new-style typed PRNG keys

Jake VanderPlas committed 2 years ago

654ae1a4

Merge pull request #3277 from IvyZX:pen

a-googler committed 2 years ago

d9f83bf7

Merge pull request #3334 from google:0.7.5

a-googler committed 2 years ago

302a7fa9

Fix issue with rng check for new-style PRNG keys

Jake VanderPlas committed 2 years ago

d6f744a4

cgarciae committed 2 years ago

33b460cc

Merge pull request #3331 from levskaya:version_fix

a-googler committed 2 years ago

9c40e030

Update minimum python version info.

levskaya committed 2 years ago

7ce981bc

Bump version no by 1 post-release.

IvyZX committed 2 years ago

ca3ea06f

Merge pull request #3326 from IvyZX:push

a-googler committed 2 years ago

5f06f08c

IvyZX committed 2 years ago

ad8d22e9

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

Older