google/flax

Pull Requests Commits

nnx: add interacting with JAX section to README

cgarciae committed 2 years ago

0a67fc77

Fix promotion bug in MultiHeadDotProductAttention:

giovannic committed 2 years ago

1c27546e

Merge pull request #3561 from chiamp:typo

a-googler committed 2 years ago

f154a64d

Merge pull request #3563 from chiamp:readme

a-googler committed 2 years ago

ca7d7483

Merge pull request #3566 from google:nnx-attention-add-qkv-features

a-googler committed 2 years ago

a59675c1

Merge pull request #3547 from google:nnx-rngs-impl-getattr

a-googler committed 2 years ago

eb1bfb11

Passed previously unused `iterate_clz` to `register_pytree_with_keys`.

superbobry committed 2 years ago

0697987e

Merge pull request #3565 from google:cgarciae-nnx-fix-readme-install

a-googler committed 2 years ago

fa2d80ef

add kwargs support for nn.jit

Cristian Garcia committed 2 years ago

cf7311da

nnx: add complex test with scan + batchnorm + dropout

cgarciae committed 2 years ago

df60f611

implement Rng.__getattr__

cgarciae committed 2 years ago

8520ad94

nnx: add qkv_features back to MHA

cgarciae committed 2 years ago

eb850d87

Merge pull request #3523 from google:improve-variables-mapping

a-googler committed 2 years ago

b4682076

[nnx] Fix readme install instruction

cgarciae committed 2 years ago

Verified be01a263

chiamp committed 2 years ago

d7eb6cdd

chiamp committed 2 years ago

64fffb01

Merge pull request #3560 from google:cgarciae-patch-1

a-googler committed 2 years ago

c9970933

Merge pull request #3556 from chiamp:readme

a-googler committed 2 years ago

0a337cbc

Merge pull request #3554 from chiamp:nnx_features

a-googler committed 2 years ago

82e1a27d

Disable ruff sort imports

cgarciae committed 2 years ago

Verified 4e852f57

updated NNX readme

chiamp committed 2 years ago

49e6f4f1

re-factored nnx initializers

chiamp committed 2 years ago

1b7493c0

re-factored features arg

chiamp committed 2 years ago

aa1703b9

Merge pull request #3553 from chiamp:format

a-googler committed 2 years ago

44a8657b

Merge pull request #3550 from chiamp:attention

a-googler committed 2 years ago

2414d1cc

format linen_linear_test.py

chiamp committed 2 years ago

71b5c464

changed to for attention layer

chiamp committed 2 years ago

7f10ad8d

Merge pull request #3548 from chiamp:nnx_attention

a-googler committed 2 years ago

2e51a4e5

Merge pull request #3537 from chiamp:nnx_conv

a-googler committed 2 years ago

73a0562e

added sow attention weights to NNX

chiamp committed 2 years ago

403ab153

Newer Older