google/flax

Pull Requests Commits

[nnx] refactor UpdateContext

cgarciae committed 1 year ago

a1e3bbf8

[nnx] remove transforms.deprecated

cgarciae committed 1 year ago

8442dc84

[nnx] support Array leaves in graph nodes

cgarciae committed 1 year ago

0ce2b0a6

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

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

Older