google/flax

Pull Requests Commits

[flax.linen] Change nn.Module.sow to enable interventions at sown values using intercept_methods.

a-googler committed 2 years ago

95a9f570

Merge pull request #3721 from google:fix-tabulate-flops

a-googler committed 2 years ago

f5c48fe8

Merge pull request #3680 from google:nnx-demo-2

a-googler committed 2 years ago

41ae533a

Merge pull request #3707 from google:nnx-simplify-readme

a-googler committed 2 years ago

9541acbd

Use `jax.tree_util.tree_map` instead of deprecated `jax.tree_map`.

IvyZX committed 2 years ago

0797d053

fix Tabulate's compute_flops

cgarciae committed 2 years ago

5ff27fdf

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

cgarciae committed 2 years ago

Verified 1fa75777

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

[nnx] simplify readme

cgarciae committed 2 years ago

26479fcc

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

Older