google/flax

Pull Requests Commits

Preserve sharding information in axes_scan

oliverdutton-iso committed 325 days ago

4e2becd8

Internal change

kho committed 326 days ago

cbb4f80f

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

a-googler committed 326 days ago

171b28e8

[nnx] allow method calls in ToLinen

cgarciae committed 326 days ago

61d0e9e0

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

a-googler committed 333 days ago

99f8e598

Merge pull request #4798 from google:ver-0.10.7

a-googler committed 333 days ago

1ecd34f9

update version to 0.10.7

cgarciae committed 333 days ago

d2cf8c2b

Fix typo in Flax `nnx_basics` doc.

a-googler committed 338 days ago

cca78723

Internal changes

basioli-k committed 338 days ago

6770f73f

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

a-googler committed 339 days ago

7d22696c

Merge pull request #4788 from google:add

a-googler committed 340 days ago

0b9bb4ee

Exposed OptState in nnx module

vfdev-5 committed 340 days ago

755067df

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

yashk2810 committed 346 days ago

dcd2f5e2

[nnx] mutable array p3

cgarciae committed 348 days ago

571c9e40

[nnx] add optax optimizer

cgarciae committed 351 days ago

0a4e8ef1

Merge pull request #4766 from IvyZX:bdg

a-googler committed 351 days ago

21e64ece

Simplify ToNNX access of Linen module methods

IvyZX committed 352 days ago

6a6561e5

Try to fix colab link for nnx docs

vfdev-5 committed 354 days ago

Verified 63a53ae4

revise axes_scan to flatten argument pytrees only once

mattjj committed 358 days ago

893a6602

Merge pull request #4770 from IvyZX:fix-links

a-googler committed 359 days ago

8233415b

Fix bad doc links

IvyZX committed 359 days ago

2d4df43e

Merge pull request #4768 from hamogu:patch-1

a-googler committed 360 days ago

c5d5ab1e

Merge pull request #4767 from nireekshak:test-branch

a-googler committed 360 days ago

2a9e5f2d

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

a-googler committed 360 days ago

f525c029

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

a-googler committed 361 days ago

096b8a02

Fix hardcoded link to filter guide in docs

hamogu committed 361 days ago

Verified ac6ca41a

Modified the Functional API link

nireekshak committed 362 days 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