google/flax

Pull Requests Commits

[nnx] add deepseek example

cgarciae committed 1 year ago

d0740d46

deepseek example

Cristian Garcia committed 1 year ago

5c7f6d0e

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

a-googler committed 1 year ago

881685c7

fix LoRA initialization error in nnx layer

a-googler committed 1 year ago

00ee6f11

Merge pull request #4503 from IvyZX:varmap

a-googler committed 1 year ago

a8a192ff

Remove all Param(None) lines

IvyZX committed 1 year ago

ad642ef6

Rename variable string mapping utils and move them to variableslib

IvyZX committed 1 year ago

61634e47

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

IvyZX committed 1 year ago

a7b0edeb

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

a-googler committed 1 year ago

d28f03fe

[nnx] fix fiddle

cgarciae committed 1 year ago

13013621

Update einsum layer for Gemma example

a-googler committed 1 year ago

0eeab5aa

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

a-googler committed 1 year ago

e3bcc442

[nnx] improve Module docs

cgarciae committed 1 year 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 1 year ago

b5d4ed85

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

a-googler committed 1 year ago

e4418e27

[nnx] fix ToNNX linen_attributes update

cgarciae committed 1 year ago

342a5e83

Merge pull request #4485 from google:nnx-multimetric

a-googler committed 1 year ago

969f3cb2

[nnx] fix MultiMetric typing

cgarciae committed 1 year ago

ca4dc8ca

Merge pull request #4479 from jakevdp:fix-workflows

a-googler committed 1 year ago

1961c125

Merge pull request #4478 from jakevdp:nightly-build

a-googler committed 1 year ago

6abf8bcb

CI: pin actions to specific commits

jakevdp committed 1 year ago

9a362592

CI: add scheduled test against nightly JAX releases

jakevdp committed 1 year ago

44efae32

[nnx] add state summaries for print and display

Cristian Garcia committed 1 year ago

3d1f55a0

Copybara import of the project:

a-googler committed 1 year ago

62688664

Merge pull request #4438 from google:nnx-tabulate

a-googler committed 1 year ago

adbad95e

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

IvyZX committed 1 year ago

9ebdbdc2

[nnx] add tabulate

cgarciae committed 1 year ago

c59bc1a9

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

a-googler committed 1 year ago

5408c40c

Makes flax Modules more compatible with IPython auto-reload.

a-googler committed 1 year ago

e2134afa

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

a-googler committed 1 year ago

dd70960a

Older