google/flax

Pull Requests Commits

Add attribute path customization to bridge modules

IvyZX committed 291 days ago

2dfd3019

Merge pull request #4527 from google:update-tracer-leakage-detection

a-googler committed 291 days ago

3e20cb9c

Merge pull request #4553 from sarlinpe:sarlinpe/conv-padding-reflect

a-googler committed 291 days ago

87e07186

Copybara import of the project:

Cristian Garcia committed 291 days ago

f14206cd

fix trace-level detection

cgarciae committed 292 days ago

b75de891

Merge pull request #4621 from google:nnx-fix-performance-docs

a-googler committed 292 days ago

96f9119c

fix performance guide

cgarciae committed 292 days ago

5b7b87f2

Merge pull request #4619 from IvyZX:bdg-name

a-googler committed 293 days ago

960b544f

support swap model params in gemma sampler

tianshub committed 293 days ago

bd9eddf2

Fixed typo in documentation for `save_checkpoint_multiprocess`.

agutkin committed 293 days ago

fefc4c9b

Fixed typo in `beam_search` loop.

agutkin committed 293 days ago

4741848c

allow bridge module to have 'name' field

IvyZX committed 294 days ago

6147d8db

Copybara import of the project:

a-googler committed 294 days ago

187b910e

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

a-googler committed 294 days ago

7c36812f

add promote_dtype as a config option for multiple layers

cgarciae committed 294 days ago

525c3d87

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

a-googler committed 295 days ago

fab37ad6

[nnx] add support for standalone Variables

cgarciae committed 295 days ago

d4aa2489

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

a-googler committed 295 days ago

0362d956

Dramatically speed up sampling compilation time

a-googler committed 295 days ago

28b5b425

Merge pull request #4604 from IvyZX:linx-nn

a-googler committed 298 days ago

d59132d9

Copybara import of the project:

a-googler committed 299 days ago

e3789de2

Add Sow Config to from_params constructor.

casaro committed 299 days ago

64635457

Fix position and name of Post Attention Norm.

casaro committed 299 days ago

b28822f3

add top_p sampling in gemma example

tianshub committed 299 days ago

a149b6d7

Add option to load checkpoints with transposed Gating Einsum.

casaro committed 299 days ago

ceb7aedd

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

a-googler committed 299 days ago

bf4c322a

Make RoPE Base Frequency configurable.

casaro committed 299 days ago

324bbe0f

Add configurable Query Pre Attention scalar.

casaro committed 299 days ago

fd32d453

bridge module with linen module

IvyZX committed 300 days ago

300494b0

Merge pull request #4584 from IvyZX:linx-misc

a-googler committed 300 days ago

8254dd0a

Older