google/flax

Pull Requests Commits

[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

[nnx] fix iter_nodes

cgarciae committed 2 years ago

90c2687b

Merge pull request #3876 from google:nnx-v0.1

a-googler committed 2 years ago

2c7d7cd9

cgarciae committed 2 years ago

08b13368

Fix jax.tree_util.register_dataclass in older JAX versions.

epiqueras committed 2 years ago

ea3bcab6

Merge pull request #3887 from chiamp:autosummary

a-googler committed 2 years ago

b8ccb153

Merge pull request #3886 from chiamp:robots

a-googler committed 2 years ago

79a0a32b

fixed autosummary links

chiamp committed 2 years ago

3b7f103f

updated robots.txt

chiamp committed 2 years ago

28caf0b5

Merge pull request #3827 from kaixih:fp8_meta_custom_dtype_non_jit

a-googler committed 2 years ago

78189324

allow custom dot_general for einsum.

a-googler committed 2 years ago

b386b715

option of forcing the input of softmax to be fp32 for better numerical stability in mixed-precision training.

a-googler committed 2 years ago

e6379c1b

Merge pull request #3871 from chiamp:nnx_api

a-googler committed 2 years ago

034b4aad

[Flax] Update JAX version check for dataclass registration.

epiqueras committed 2 years ago

33911e60

Add builtin cc dataclass pytree node for performance.

epiqueras committed 2 years ago

9b512989

Newer Older