google/flax

Pull Requests Commits

[nnx] add Module.move

cgarciae committed 1 year ago

07877ad9

[nnx] add pytree support

cgarciae committed 1 year ago

9383573c

[nnx] cleanup graph_utils

cgarciae committed 1 year ago

f3b57e18

Merge pull request #3720 from google:nnx-explicit-variables

a-googler committed 1 year ago

7d6de00c

[nnx] explicit Variables

cgarciae committed 1 year ago

0d347d03

Merge pull request #3722 from google:nnx-dynamic-rngs

a-googler committed 1 year ago

acba0bf6

Merge pull request #3704 from google:nnx-fix-state-sub

a-googler committed 1 year ago

a1ecdf73

[nnx] dynamic RngStream.count + ForkRngs

cgarciae committed 1 year ago

3e693506

Merge pull request #3731 from IvyZX:conv-trans

a-googler committed 1 year ago

1abfa872

Add broadcast of strides and kernel_dilation to ConvTranspose

IvyZX committed 1 year ago

fdbc640c

[nnx] Fix State.__sub__

cgarciae committed 1 year ago

4dcc0ef4

Merge pull request #3724 from google:nnx-simplify-trace-state

a-googler committed 1 year ago

d1f219ff

[nnx] simplify TraceState

cgarciae committed 1 year ago

4b7ba096

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

Older