google/flax

Pull Requests Commits

Merge pull request #4813 from google:to-linen-eval-shape-2

a-googler committed 1 year ago

8c22f295

modify NNX to use id(variable) instead of nnx.Variables as dictionary keys

divyashreepathihalli committed 1 year ago

f411a322

Merge pull request #4776 from google:add-note-gemma-guide-about-model-license

a-googler committed 1 year ago

23d1a1b4

Merge pull request #4790 from google:enable-docs-nnx-remat

a-googler committed 1 year ago

a3f5be21

[nnx] ToLinen add abtract_init flag

cgarciae committed 1 year ago

7268e5f5

Merge pull request #4809 from IvyZX:doc2

a-googler committed 1 year ago

188baabc

Deduplicate contributing and philosophy and move to main site

IvyZX committed 1 year ago

103f65de

Preserve sharding information in axes_scan

oliverdutton-iso committed 1 year ago

4e2becd8

Internal change

kho committed 1 year ago

cbb4f80f

Merge pull request #4808 from google:to-linen-method-calls

a-googler committed 1 year ago

171b28e8

[nnx] allow method calls in ToLinen

cgarciae committed 1 year ago

61d0e9e0

[nnx] remove VariableState

cgarciae committed 1 year ago

e8984b70

Merge pull request #4755 from google:nnx-remove-all

a-googler committed 1 year ago

99f8e598

Merge pull request #4798 from google:ver-0.10.7

a-googler committed 1 year ago

1ecd34f9

update version to 0.10.7

cgarciae committed 1 year ago

d2cf8c2b

Fix typo in Flax `nnx_basics` doc.

a-googler committed 1 year ago

cca78723

Internal changes

basioli-k committed 1 year ago

6770f73f

Merge pull request #4775 from google:try-fix-colab-link-nnx

a-googler committed 1 year ago

7d22696c

Fixed nnx.remat docstring rendering

vfdev-5 committed 1 year ago

Verified 4870afe5

Merge pull request #4788 from google:add

a-googler committed 1 year ago

0b9bb4ee

Exposed OptState in nnx module

vfdev-5 committed 1 year ago

755067df

Use `.input_formats` and `.output_formats` in place of `.input_layouts` and `.output_layouts` respectively.

yashk2810 committed 1 year ago

dcd2f5e2

[nnx] mutable array p3

cgarciae committed 1 year ago

571c9e40

[nnx] add optax optimizer

cgarciae committed 1 year ago

0a4e8ef1

Merge pull request #4766 from IvyZX:bdg

a-googler committed 1 year ago

21e64ece

Simplify ToNNX access of Linen module methods

IvyZX committed 1 year ago

6a6561e5

vfdev-5 committed 1 year ago

Verified 62238aaf

Added a note to gemma guide about model's license consent on kaggle

vfdev-5 committed 1 year ago

Verified 3945f066

Try to fix colab link for nnx docs

vfdev-5 committed 1 year ago

Verified 63a53ae4

revise axes_scan to flatten argument pytrees only once

mattjj committed 1 year ago

893a6602

Newer Older