google/flax

Pull Requests Commits

[nnx] diffusion example

cgarciae committed 1 year ago

b41bb232

[nnx] Sequential uses regular list

cgarciae committed 1 year ago

3ab233b3

[nnx] add ConvTranspose

cgarciae committed 1 year ago

a11f6f28

Merge pull request #3901 from NeilGirdhar:fix_threading

a-googler committed 1 year ago

40c7eaa7

Share nnx node registry between threads

NeilGirdhar committed 1 year ago

ba935a0d

Copybara import of the project:

tkoeppe committed 1 year ago

859435e4

Merge pull request #3880 from NeilGirdhar:fix_mappings

a-googler committed 1 year ago

d0e080da

Make FlatState a Mapping instead of a dict

NeilGirdhar committed 1 year ago

d1935c71

Merge pull request #3895 from chiamp:codecov

a-googler committed 1 year ago

ae5d66d3

chiamp committed 1 year ago

b4b9d12e

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

a-googler committed 1 year ago

2c7d7cd9

cgarciae committed 1 year ago

08b13368

Fix jax.tree_util.register_dataclass in older JAX versions.

epiqueras committed 1 year ago

ea3bcab6

Merge pull request #3887 from chiamp:autosummary

a-googler committed 1 year ago

b8ccb153

Merge pull request #3886 from chiamp:robots

a-googler committed 1 year ago

79a0a32b

fixed autosummary links

chiamp committed 1 year ago

3b7f103f

updated robots.txt

chiamp committed 1 year ago

28caf0b5

Merge pull request #3827 from kaixih:fp8_meta_custom_dtype_non_jit

a-googler committed 1 year ago

78189324

allow custom dot_general for einsum.

a-googler committed 1 year ago

b386b715

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

a-googler committed 1 year ago

e6379c1b

Merge pull request #3871 from chiamp:nnx_api

a-googler committed 1 year ago

034b4aad

[Flax] Update JAX version check for dataclass registration.

epiqueras committed 1 year ago

33911e60

Add builtin cc dataclass pytree node for performance.

epiqueras committed 1 year ago

9b512989

Work with older jax versions

kaixih committed 1 year ago

5dfc0dbc

Remove `ocdbt_merge` and `restore_with_serialized_types` options from `PyTreeCheckpointHandler` as they are not in use anywhere.

cpgaffney1 committed 1 year ago

fa119fda

Add option to skip float32 promotion when computing means and variances for normalization.

a-googler committed 1 year ago

0310ac7d

Fix LogicalRules type annotation. (Tuple[str] is a tuple with single element string,

marksandler2 committed 1 year ago

b473704c

Merge pull request #3841 from chiamp:dropout2

a-googler committed 1 year ago

e674759f

updated nnx dropout

chiamp committed 1 year ago

787c96f7

added nnx api reference link

chiamp committed 1 year ago

82aa6595

Older