google/flax

Pull Requests Commits

[linen] fix DenseGeneral init

Cristian Garcia committed 1 year ago

c6f097d4

Merge pull request #3835 from google:linen-attention-multiple-initializers

a-googler committed 1 year ago

e0fa96fe

Merge pull request #3891 from google:v0.8.4

a-googler committed 1 year ago

ee8c55e4

Merge pull request #3900 from google:linen-test-jit-state

a-googler committed 1 year ago

811098ec

Merge pull request #3913 from chiamp:path

a-googler committed 2 years ago

f02c5d51

cgarciae committed 2 years ago

3744eefb

Merge pull request #3864 from google:nnx-static-issue

a-googler committed 2 years ago

4ba62991

Merge pull request #3909 from google:nnx-sequential-uses-list

a-googler committed 2 years ago

51f54638

Merge pull request #3889 from google:nnx-fix-iter-nodes

a-googler committed 2 years ago

a6807676

[nnx] add Module pytree_experimental static test

cgarciae committed 2 years ago

2cdd670e

added docstring for Module.scope.path

chiamp committed 2 years ago

fdc893a9

Copybara import of the project:

Cristian Garcia committed 2 years ago

b000c756

Merge pull request #3894 from google:nnx-enable-mypy-pytype-ci

a-googler committed 2 years ago

2c833a97

[nnx] Sequential uses regular list

cgarciae committed 2 years ago

3ab233b3

[nnx] add ConvTranspose

cgarciae committed 2 years ago

a11f6f28

Add support for `jax.sharding.PartitionSpec.UNCONSTRAINED` in logical specification

marksandler2 committed 2 years ago

d237a45f

Merge pull request #3847 from chiamp:haiku

a-googler committed 2 years ago

ea025eb1

updated haiku guide

chiamp committed 2 years ago

59e037f0

Merge pull request #3905 from chiamp:lm1b

a-googler committed 2 years ago

3d98edad

fixed jnp.clip deprecation

chiamp committed 2 years ago

fb5515c1

[nnx] enable pytype on CI

cgarciae committed 2 years ago

0459456f

Merge pull request #3901 from NeilGirdhar:fix_threading

a-googler committed 2 years ago

40c7eaa7

Share nnx node registry between threads

NeilGirdhar committed 2 years ago

ba935a0d

[linen] test jit caching with state updates

cgarciae committed 2 years ago

e70f6bbe

Copybara import of the project:

tkoeppe committed 2 years ago

859435e4

Merge pull request #3880 from NeilGirdhar:fix_mappings

a-googler committed 2 years ago

d0e080da

Make FlatState a Mapping instead of a dict

NeilGirdhar committed 2 years ago

d1935c71

[nnx] enable pytype and mypy on CI

cgarciae committed 2 years ago

895b7833

Merge pull request #3895 from chiamp:codecov

a-googler committed 2 years ago

ae5d66d3

chiamp committed 2 years ago

b4b9d12e

Older