google/flax

Pull Requests Commits

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

a-googler committed 1 year ago

95a9f570

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

a-googler committed 1 year ago

f5c48fe8

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

a-googler committed 1 year ago

41ae533a

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

a-googler committed 1 year ago

9541acbd

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

IvyZX committed 1 year ago

0797d053

fix Tabulate's compute_flops

cgarciae committed 1 year ago

5ff27fdf

Replace deprecated API `jax.tree_map`

JW1992 committed 1 year ago

83e74665

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

cgarciae committed 1 year ago

Verified 1fa75777

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

[nnx] simplify readme

cgarciae committed 1 year ago

26479fcc

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

Older