google/flax

Pull Requests Commits

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

Copybara import of the project:

Cristian Garcia committed 1 year ago

1fbe17d8

Merge pull request #3861 from gnecula:nnx_doc1

a-googler committed 1 year ago

99736930

integrate changes to nnx_basics

cgarciae committed 1 year ago

7f1b36a7

Merge pull request #3099 from google:rnncell-refactor-flip

a-googler committed 1 year ago

bf4dffff

[nnx] Some small documentation suggestions.

gnecula committed 1 year ago

05a72fc6

Merge pull request #3859 from google:nnx-improve-nnx-basics

a-googler committed 1 year ago

f7c0669d

Merge pull request #3843 from IvyZX:jit-guide

a-googler committed 1 year ago

9c008e33

[nnx] improve nnx basics

cgarciae committed 1 year ago

1f27faa8

Copybara import of the project:

Cristian Garcia committed 1 year ago

3711fa60

Merge pull request #3817 from google:nnx-jit-apply-sharding

a-googler committed 1 year ago

2438518b

Older