google/flax

Pull Requests Commits

Merge pull request #3927 from google:nnx-cond

a-googler committed 1 year ago

7542b28a

[nnx] transforms refactor

cgarciae committed 1 year ago

6b31a980

Merge pull request #3921 from google:nnx-compact

a-googler committed 1 year ago

2bdcee76

[nnx] add compat

cgarciae committed 1 year ago

72ea9599

Merge pull request #3919 from google:nnx-iternodes-docs

a-googler committed 1 year ago

f2b2ea33

[nnx] add iter_graph

cgarciae committed 1 year ago

ed49ab36

Merge pull request #3910 from google:nnx-object-refactor

a-googler committed 1 year ago

cf51da39

Added a missing return to _fingerprint_recursive

superbobry committed 1 year ago

6b6014b0

Replace deprecated `jax.tree_*` functions with `jax.tree.*`

Jake VanderPlas committed 1 year ago

d37d8349

[nnx] Object refactor

cgarciae committed 1 year ago

24002bc7

Merge pull request #3925 from google:nnx-fix-bugs

a-googler committed 1 year ago

e67c9716

cgarciae committed 1 year ago

efc912a1

Merge pull request #3915 from google:nnx-graph-comments

a-googler committed 1 year ago

a2a62425

[nnx] cleanup graph

cgarciae committed 1 year ago

7b611e37

[nnx] enable pytype and mypy on CI

Cristian Garcia committed 1 year ago

8452b4e6

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 1 year ago

f02c5d51

cgarciae committed 1 year ago

3744eefb

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

a-googler committed 1 year ago

4ba62991

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

a-googler committed 1 year ago

51f54638

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

a-googler committed 1 year ago

a6807676

[nnx] add Module pytree_experimental static test

cgarciae committed 1 year ago

2cdd670e

added docstring for Module.scope.path

chiamp committed 1 year ago

fdc893a9

Copybara import of the project:

Cristian Garcia committed 1 year ago

b000c756

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

a-googler committed 1 year ago

2c833a97

[nnx] Sequential uses regular list

cgarciae committed 1 year ago

3ab233b3

[nnx] add ConvTranspose

cgarciae committed 1 year ago

a11f6f28

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

marksandler2 committed 1 year ago

d237a45f

Newer Older