google/flax

Pull Requests Commits

Merge pull request #3737 from google:jit-fingerprint

a-googler committed 2 years ago

6760e32c

nn.jit: automatic fingerprint definition for dataclass attributes

cgarciae committed 2 years ago

3ea30d71

Merge pull request #3710 from chiamp:einsum

a-googler committed 2 years ago

49a8dde5

added Einsum layer

chiamp committed 2 years ago

8e564422

Flax: avoid key reuse in tests

Jake VanderPlas committed 2 years ago

a307d473

Improves fingerprint definition for Modules in nn.jit.

Cristian Garcia committed 2 years ago

f195a8b4

[nnx] cleanup graph_utils

cgarciae committed 2 years ago

f3b57e18

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

a-googler committed 2 years ago

7d6de00c

[nnx] explicit Variables

cgarciae committed 2 years ago

0d347d03

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

a-googler committed 2 years ago

acba0bf6

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

a-googler committed 2 years ago

a1ecdf73

[nnx] dynamic RngStream.count + ForkRngs

cgarciae committed 2 years ago

3e693506

Merge pull request #3731 from IvyZX:conv-trans

a-googler committed 2 years ago

1abfa872

Add broadcast of strides and kernel_dilation to ConvTranspose

IvyZX committed 2 years ago

fdbc640c

[nnx] Fix State.__sub__

cgarciae committed 2 years ago

4dcc0ef4

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

a-googler committed 2 years ago

d1f219ff

[nnx] simplify TraceState

cgarciae committed 2 years ago

4b7ba096

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

kaixih committed 2 years ago

dd004c21

Add the full attr

kaixih committed 2 years ago

54fa392b

Use custom grad accum for fp8 meta params

kaixih committed 2 years ago

02913030

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

Newer Older