google/flax

Pull Requests Commits

[nnx] improve Variable + graph_utils

cgarciae committed 1 year ago

c814f20f

[nnx] explicit Variables

cgarciae committed 1 year ago

5f5dd381

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

a-googler committed 1 year ago

aaf130c9

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

a-googler committed 1 year ago

746fc65d

fix Module.module_paths docstring

cgarciae committed 1 year ago

17fbdff8

Merge pull request #3697 from carlosgmartin:simple_rnn_cell

a-googler committed 1 year ago

b3ac6cbd

Add SimpleCell.

carlosgmartin committed 1 year ago

6be118c3

Merge pull request #3699 from chiamp:make_rng

a-googler committed 1 year ago

c83ee86d

make make_rng default to 'params'

chiamp committed 1 year ago

c6b5de01

use the latest jax on CI

cgarciae committed 1 year ago

6a9a98fd

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

a-googler committed 1 year ago

736d1331

Merge pull request #3698 from chiamp:apply_rng

a-googler committed 1 year ago

daf06ead

Merge pull request #3703 from levskaya:rematdocfix

a-googler committed 1 year ago

eca5dd6c

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

levskaya committed 1 year ago

5c7fb4c9

added default params rng to .apply

chiamp committed 1 year ago

925fe9c0

Merge pull request #3678 from google:nnx-quickstart2

a-googler committed 1 year ago

efbe7050

make perturb() work in setup mode

levskaya committed 1 year ago

61812398

Merge pull request #3694 from chiamp:intercept_methods

a-googler committed 1 year ago

e025e593

Merge pull request #3693 from chiamp:nnx_compute_stats

a-googler committed 1 year ago

d061da84

fixed intercept_methods docstring

chiamp committed 1 year ago

92532f1a

updated nnx compute_stats

chiamp committed 1 year ago

a4511079

Merge pull request #3691 from chiamp:nnx_rms

a-googler committed 1 year ago

c25f546f

Fix bug that assumed frozen-dict keys were strings.

levskaya committed 1 year ago

ce156783

added nnx rmsnorm

chiamp committed 1 year ago

172078e9

Merge pull request #3687 from chiamp:nnx_attention

a-googler committed 1 year ago

a5661961

added nnx attention equivalence test

chiamp committed 1 year ago

6956a2d5

[nnx] add Quick Start

cgarciae committed 1 year ago

7837c591

Merge pull request #3688 from chiamp:init_cache

a-googler committed 1 year ago

c2c81ee1

Merge pull request #3689 from chiamp:nnx_norm

a-googler committed 1 year ago

7330d6fc

added nnx normalization tests

chiamp committed 1 year ago

fa324c0b

Older