google/flax

Pull Requests Commits

fixed intercept_methods docstring

chiamp committed 2 years ago

fca6b454

removed getattr/setattr magic from nnx.Pytree

chiamp committed 2 years ago

c279f248

[nnx] add pytree support

cgarciae committed 2 years ago

c8987746

Merge pull request #3753 from IvyZX:orbax-quickfix

a-googler committed 2 years ago

ce8a3c74

Fix doctest and unbreak head

IvyZX committed 2 years ago

3b23dc8c

Fix instances of JAX key reuse in flax tests

Jake VanderPlas committed 2 years ago

6f2b08e0

Merge pull request #3728 from google:nnx-cleanup-graph-utils

a-googler committed 2 years ago

0c7d256d

Merge pull request #3742 from chiamp:nnx_dataclass

a-googler committed 2 years ago

e8e0b2df

removed nnx dataclass

chiamp committed 2 years ago

d4ac6c3f

Merge pull request #3623 from kaixih:fp8_meta_custom_grad_accumulation

a-googler committed 2 years ago

e4282ee1

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

Newer Older