google/flax

Pull Requests Commits

deepseek example

Cristian Garcia committed 331 days ago

5c7f6d0e

Merge pull request #4504 from IvyZX:no-none-var

a-googler committed 331 days ago

881685c7

fix LoRA initialization error in nnx layer

a-googler committed 332 days ago

00ee6f11

Merge pull request #4503 from IvyZX:varmap

a-googler committed 334 days ago

a8a192ff

Remove all Param(None) lines

IvyZX committed 335 days ago

ad642ef6

Rename variable string mapping utils and move them to variableslib

IvyZX committed 335 days ago

61634e47

Don't create param in normalization layers instead of create None-value params.

IvyZX committed 335 days ago

a7b0edeb

Merge pull request #4500 from google:nnx-fix-fiddle

a-googler committed 336 days ago

d28f03fe

[nnx] fix fiddle

cgarciae committed 337 days ago

13013621

Update einsum layer for Gemma example

a-googler committed 337 days ago

0eeab5aa

Merge pull request #4499 from google:nnx-improve-module-docs

a-googler committed 337 days ago

e3bcc442

[nnx] improve Module docs

cgarciae committed 337 days ago

bdcc33a5

Remove usages of orbax_utils.save_args_from_target, as this function does nothing (it used to control a checkpointing behavior that has since been optimized away).

cpgaffney1 committed 337 days ago

b5d4ed85

Merge pull request #4486 from google:fix-to-nnx-linen-attributes

a-googler committed 344 days ago

e4418e27

[nnx] fix ToNNX linen_attributes update

cgarciae committed 344 days ago

342a5e83

Merge pull request #4485 from google:nnx-multimetric

a-googler committed 344 days ago

969f3cb2

[nnx] fix MultiMetric typing

cgarciae committed 344 days ago

ca4dc8ca

Merge pull request #4479 from jakevdp:fix-workflows

a-googler committed 349 days ago

1961c125

Merge pull request #4478 from jakevdp:nightly-build

a-googler committed 349 days ago

6abf8bcb

CI: pin actions to specific commits

jakevdp committed 349 days ago

9a362592

CI: add scheduled test against nightly JAX releases

jakevdp committed 349 days ago

44efae32

[nnx] add state summaries for print and display

Cristian Garcia committed 350 days ago

3d1f55a0

Copybara import of the project:

a-googler committed 350 days ago

62688664

Merge pull request #4438 from google:nnx-tabulate

a-googler committed 350 days ago

adbad95e

Allow `nnx.bridge.variables.nnx_attrs_to_linen_vars` take `nnx.VariableState` as argument.

IvyZX committed 350 days ago

9ebdbdc2

[nnx] add tabulate

cgarciae committed 350 days ago

c59bc1a9

Merge pull request #4407 from google:rnn-broadcast-rngs

a-googler committed 351 days ago

5408c40c

Makes flax Modules more compatible with IPython auto-reload.

a-googler committed 353 days ago

e2134afa

Merge pull request #4199 from 8bitmp3:upgrade-nnx-filters

a-googler committed 353 days ago

dd70960a

Merge pull request #4472 from google:nnx-fix-fori

a-googler committed 353 days ago

595e711f

Older