google/flax

Pull Requests Commits

Use .format in place of .layout.

yashk2810 committed 346 days ago

e57eadce

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

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

Fix head CI tests

IvyZX committed 1 year ago

63fa7de7

Merge pull request #4763 from nireekshak:testing-branch1

a-googler committed 1 year ago

e4ab8833

Fix broken link for Transforms guide

nireekshak committed 1 year ago

33caa7ec

Update docs_nnx/guides/transforms.ipynb

vfdev-5 committed 1 year ago

Verified dd9096bf

Update docs_nnx/guides/transforms.md

vfdev-5 committed 1 year ago

Verified 5169897b

nireekshak committed 1 year ago

f4563283

Add in_kv_features argument to nnx.MultiHeadAttention, addressing https://github.com/google/flax/issues/4756.

Ryan McKenna committed 1 year ago

11643b4b

Fixed lm1b_nnx example training script

vfdev-5 committed 1 year ago

7def2279

Merge pull request #4741 from google:mutable-array-p2

a-googler committed 1 year ago

a5eebe5d

MultiHeadAttention only keeps rngs if dropout_rate is positive

Cristian Garcia committed 1 year ago

65dfd465

move freeze, mutable, and pure to graph.py

cgarciae committed 1 year ago

f50b4d98

Merge pull request #4749 from vfdev-5:mha-add-keep-rngs-arg

a-googler committed 1 year ago

45ce6ee2

Merge pull request #4747 from vfdev-5:fix-docs-eval_shape

a-googler committed 1 year ago

f265d914

Merge pull request #4746 from vfdev-5:update-pre-commit-hooks-version

a-googler committed 1 year ago

172ac348

[flax] unconditionally register nnx.Variable as a pytree

Jake VanderPlas committed 1 year ago

866fe952

Older