google/flax

Pull Requests Commits

Preserve sharding information in axes_scan

oliverdutton-iso committed 348 days ago

4e2becd8

Internal change

kho committed 349 days ago

cbb4f80f

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

a-googler committed 349 days ago

171b28e8

[nnx] allow method calls in ToLinen

cgarciae committed 349 days ago

61d0e9e0

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

a-googler committed 356 days ago

99f8e598

Merge pull request #4798 from google:ver-0.10.7

a-googler committed 356 days ago

1ecd34f9

update version to 0.10.7

cgarciae committed 356 days ago

d2cf8c2b

Fix typo in Flax `nnx_basics` doc.

a-googler committed 361 days ago

cca78723

Internal changes

basioli-k committed 361 days ago

6770f73f

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

a-googler committed 362 days ago

7d22696c

Merge pull request #4788 from google:add

a-googler committed 363 days ago

0b9bb4ee

Exposed OptState in nnx module

vfdev-5 committed 363 days 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

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

Merge pull request #4770 from IvyZX:fix-links

a-googler committed 1 year ago

8233415b

Fix bad doc links

IvyZX committed 1 year ago

2d4df43e

Merge pull request #4768 from hamogu:patch-1

a-googler committed 1 year ago

c5d5ab1e

Merge pull request #4767 from nireekshak:test-branch

a-googler committed 1 year ago

2a9e5f2d

Merge pull request #4563 from muhrin:fix-1925-linen-nonstring-keys

a-googler committed 1 year ago

f525c029

Check for leaves of type variablelib.Variable when getting sharding specs.

a-googler committed 1 year ago

096b8a02

Fix hardcoded link to filter guide in docs

hamogu committed 1 year ago

Verified ac6ca41a

Modified the Functional API link

nireekshak committed 1 year ago

360d8141

Merge pull request #4725 from nireekshak:testing-branch

a-googler committed 1 year ago

f88d423c

Merge pull request #4764 from IvyZX:ci

a-googler committed 1 year ago

33927d04

Merge pull request #4745 from vfdev-5:fix-lm1b-nnx-example

a-googler committed 1 year ago

a2fe3ca2

Older