google/flax

Pull Requests Commits

Replace deprecated API `jax.tree_map`

JW1992 committed 2 years ago

83e74665

Merge pull request #3705 from google:use-latest-jax

a-googler committed 2 years ago

aaf130c9

Merge pull request #3709 from google:fix-module-paths-docs

a-googler committed 2 years ago

746fc65d

fix Module.module_paths docstring

cgarciae committed 2 years ago

17fbdff8

Merge pull request #3697 from carlosgmartin:simple_rnn_cell

a-googler committed 2 years ago

b3ac6cbd

Add SimpleCell.

carlosgmartin committed 2 years ago

6be118c3

Merge pull request #3699 from chiamp:make_rng

a-googler committed 2 years ago

c83ee86d

make make_rng default to 'params'

chiamp committed 2 years ago

c6b5de01

use the latest jax on CI

cgarciae committed 2 years ago

6a9a98fd

Merge pull request #3674 from google:nnx-partial-init

a-googler committed 2 years ago

736d1331

Merge pull request #3698 from chiamp:apply_rng

a-googler committed 2 years ago

daf06ead

Merge pull request #3703 from levskaya:rematdocfix

a-googler committed 2 years ago

eca5dd6c

Fix pointless docstring example of nn.checkpoint / nn.remat.

levskaya committed 2 years ago

5c7fb4c9

added default params rng to .apply

chiamp committed 2 years ago

925fe9c0

Merge pull request #3678 from google:nnx-quickstart2

a-googler committed 2 years ago

efbe7050

make perturb() work in setup mode

levskaya committed 2 years ago

61812398

Merge pull request #3694 from chiamp:intercept_methods

a-googler committed 2 years ago

e025e593

Merge pull request #3693 from chiamp:nnx_compute_stats

a-googler committed 2 years ago

d061da84

fixed intercept_methods docstring

chiamp committed 2 years ago

92532f1a

updated nnx compute_stats

chiamp committed 2 years ago

a4511079

Merge pull request #3691 from chiamp:nnx_rms

a-googler committed 2 years ago

c25f546f

Fix bug that assumed frozen-dict keys were strings.

levskaya committed 2 years ago

ce156783

added nnx rmsnorm

chiamp committed 2 years ago

172078e9

Merge pull request #3687 from chiamp:nnx_attention

a-googler committed 2 years ago

a5661961

added nnx attention equivalence test

chiamp committed 2 years ago

6956a2d5

[nnx] add Quick Start

cgarciae committed 2 years ago

7837c591

Merge pull request #3688 from chiamp:init_cache

a-googler committed 2 years ago

c2c81ee1

Merge pull request #3689 from chiamp:nnx_norm

a-googler committed 2 years ago

7330d6fc

added nnx normalization tests

chiamp committed 2 years ago

fa324c0b

added init_cache docstring example

chiamp committed 2 years ago

c257e9e1

Older