google/flax

Pull Requests Commits

Dramatically speed up sampling compilation time

a-googler committed 1 year ago

28b5b425

Merge pull request #4604 from IvyZX:linx-nn

a-googler committed 1 year ago

d59132d9

Copybara import of the project:

a-googler committed 1 year ago

e3789de2

Add Sow Config to from_params constructor.

casaro committed 1 year ago

64635457

Fix position and name of Post Attention Norm.

casaro committed 1 year ago

b28822f3

add top_p sampling in gemma example

tianshub committed 1 year ago

a149b6d7

Add option to load checkpoints with transposed Gating Einsum.

casaro committed 1 year ago

ceb7aedd

Merge pull request #4547 from google:nnx-index-pytrees

a-googler committed 1 year ago

bf4c322a

Make RoPE Base Frequency configurable.

casaro committed 1 year ago

324bbe0f

Add configurable Query Pre Attention scalar.

casaro committed 1 year ago

fd32d453

bridge module with linen module

IvyZX committed 1 year ago

300494b0

Merge pull request #4584 from IvyZX:linx-misc

a-googler committed 1 year ago

8254dd0a

bridge module with nnx submodules

IvyZX committed 1 year ago

06a50e88

casaro committed 1 year ago

0769411f

Merge pull request #4600 from google:remove-embed-special-case

a-googler committed 1 year ago

50b996d2

remove Embed nan casting

cgarciae committed 1 year ago

0933c2e8

Merge pull request #4592 from jakevdp:fix-shape

a-googler committed 1 year ago

45a8f846

Avoid calls to jnp.shape for non-array inputs.

jakevdp committed 1 year ago

4e300d4c

[nnx] improve grad docs

cgarciae committed 1 year ago

33d34b43

[bridge] improve sow and add initializers

Cristian Garcia committed 1 year ago

a24d790e

Refactor bridge.Module tests from `wrappers_test.py` to another file.

IvyZX committed 1 year ago

6af8fcb0

Merge pull request #4580 from google:nnx-tabulate-issue

a-googler committed 1 year ago

7974b106

Merge pull request #4579 from google:v0.10.4

a-googler committed 1 year ago

7f175f80

[nnx] fix tabulate

cgarciae committed 1 year ago

14fda9d7

cgarciae committed 1 year ago

5f4f1b84

[bridge] Set _initializing correctly and avoid return RNG states

IvyZX committed 1 year ago

54138501

Merge pull request #4576 from IvyZX:fix-head

a-googler committed 1 year ago

d96be6ce

Fix breakages from newest jax

IvyZX committed 1 year ago

0f9f55d6

Merge pull request #4490 from google:nnx-shard-map

a-googler committed 1 year ago

41bef07c

Add REFLECT padding to convolution layer

sarlinpe committed 1 year ago

d552d4b8

Newer Older