google/flax

Pull Requests Commits

Fix flaky vmap test tolerance.

levskaya committed 1 year ago

0dec8a80

[nnx] remove transforms.deprecated

cgarciae committed 1 year ago

ce48e5d9

Merge pull request #4649 from IvyZX:attn

a-googler committed 1 year ago

527cb4a4

Use jax.nn.dot_product_attention when possible

IvyZX committed 1 year ago

6f7322db

Merge pull request #4648 from IvyZX:fix-head

a-googler committed 1 year ago

5f7af869

roll up package jax version and uv.lock

IvyZX committed 1 year ago

6a77c8a9

Merge pull request #4612 from google:nnx-array-leaves

a-googler committed 1 year ago

097fffce

Merge pull request #4646 from google/release-webhook

ZacCranko committed 1 year ago

Verified ba59c336

[nnx] support Array leaves in graph nodes

cgarciae committed 1 year ago

dfd1788c

discord release webhook

cgarciae committed 1 year ago

9e4d3e3f

Stop passing reduce_axes to jax.grad, jax.vjp, and jax.value_and_grad.

jburnim committed 1 year ago

f18baf07

Fix module stack typing annotation.

a-googler committed 1 year ago

9f5d1309

Merge pull request #4641 from google:nnx-improve-variable

a-googler committed 1 year ago

a9ef59b2

[nnx] improve Variable proxy for binary operations

cgarciae committed 1 year ago

6a2b33e2

Merge pull request #4634 from IvyZX:parentattr

a-googler committed 1 year ago

fa0f3e8e

Allow name arg to represent actual submodule path while keeping other attrs

IvyZX committed 1 year ago

18faf549

Merge pull request #4637 from IvyZX:custom-meta

a-googler committed 1 year ago

1e483800

allow custom axis metadata

IvyZX committed 1 year ago

a920dd22

Merge pull request #4632 from google:add-max-repr-depth

a-googler committed 1 year ago

dd3c18be

add reprlib max depth flag

cgarciae committed 1 year ago

1bdbab33

Merge pull request #4624 from IvyZX:attr-priority

a-googler committed 1 year ago

42098ad6

Add attribute path customization to bridge modules

IvyZX committed 1 year ago

c3230c02

Merge pull request #4527 from google:update-tracer-leakage-detection

a-googler committed 1 year ago

3e20cb9c

Merge pull request #4553 from sarlinpe:sarlinpe/conv-padding-reflect

a-googler committed 1 year ago

87e07186

Copybara import of the project:

Cristian Garcia committed 1 year ago

f14206cd

fix trace-level detection

cgarciae committed 1 year ago

b75de891

Merge pull request #4621 from google:nnx-fix-performance-docs

a-googler committed 1 year ago

96f9119c

fix performance guide

cgarciae committed 1 year ago

5b7b87f2

Merge pull request #4619 from IvyZX:bdg-name

a-googler committed 1 year ago

960b544f

support swap model params in gemma sampler

tianshub committed 1 year ago

bd9eddf2

Newer Older