google/flax

Pull Requests Commits

Fixed typo in documentation for `save_checkpoint_multiprocess`.

agutkin committed 1 year ago

fefc4c9b

Fixed typo in `beam_search` loop.

agutkin committed 1 year ago

4741848c

allow bridge module to have 'name' field

IvyZX committed 1 year ago

6147d8db

Copybara import of the project:

a-googler committed 1 year ago

187b910e

Merge pull request #4613 from google:fix-promote-dtypes

a-googler committed 1 year ago

7c36812f

add promote_dtype as a config option for multiple layers

cgarciae committed 1 year ago

525c3d87

Merge pull request #4606 from google:nnx-standalone-variables

a-googler committed 1 year ago

fab37ad6

[nnx] add support for standalone Variables

cgarciae committed 1 year ago

d4aa2489

Merge pull request #4588 from google:nnx-graph-docs

a-googler committed 1 year ago

0362d956

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

Newer Older