google/flax

Pull Requests Commits

Dramatically speed up sampling compilation time

a-googler committed 324 days ago

28b5b425

Merge pull request #4604 from IvyZX:linx-nn

a-googler committed 326 days ago

d59132d9

Copybara import of the project:

a-googler committed 327 days ago

e3789de2

Add Sow Config to from_params constructor.

casaro committed 328 days ago

64635457

Fix position and name of Post Attention Norm.

casaro committed 328 days ago

b28822f3

add top_p sampling in gemma example

tianshub committed 328 days ago

a149b6d7

Add option to load checkpoints with transposed Gating Einsum.

casaro committed 328 days ago

ceb7aedd

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

a-googler committed 328 days ago

bf4c322a

Make RoPE Base Frequency configurable.

casaro committed 328 days ago

324bbe0f

Add configurable Query Pre Attention scalar.

casaro committed 328 days ago

fd32d453

bridge module with linen module

IvyZX committed 328 days ago

300494b0

Merge pull request #4584 from IvyZX:linx-misc

a-googler committed 328 days ago

8254dd0a

bridge module with nnx submodules

IvyZX committed 328 days ago

06a50e88

casaro committed 329 days ago

0769411f

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

a-googler committed 329 days ago

50b996d2

remove Embed nan casting

cgarciae committed 329 days ago

0933c2e8

Merge pull request #4592 from jakevdp:fix-shape

a-googler committed 329 days ago

45a8f846

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

jakevdp committed 330 days ago

4e300d4c

[nnx] improve grad docs

cgarciae committed 333 days ago

33d34b43

[bridge] improve sow and add initializers

Cristian Garcia committed 333 days ago

a24d790e

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

IvyZX committed 335 days ago

6af8fcb0

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

a-googler committed 335 days ago

7974b106

Merge pull request #4579 from google:v0.10.4

a-googler committed 335 days ago

7f175f80

[nnx] fix tabulate

cgarciae committed 335 days ago

14fda9d7

cgarciae committed 335 days ago

5f4f1b84

[bridge] Set _initializing correctly and avoid return RNG states

IvyZX committed 336 days ago

54138501

Merge pull request #4576 from IvyZX:fix-head

a-googler committed 336 days ago

d96be6ce

Fix breakages from newest jax

IvyZX committed 336 days ago

0f9f55d6

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

a-googler committed 337 days ago

41bef07c

Add REFLECT padding to convolution layer

sarlinpe committed 339 days ago

d552d4b8

Newer Older