google/flax

Pull Requests Commits

Merge pull request #4891 from mohsinm-dev:fix-nnx-tabulate-empty-dict-issue-4889

a-googler committed 250 days ago

f65afdb8

Merge pull request #4948 from samanklesaria:issues/4883

a-googler committed 250 days ago

ebb980a6

clean up jax.Ref-related names

Cristian Garcia committed 250 days ago

4c4bc3b1

Merge pull request #4980 from IvyZX:shardref

a-googler committed 253 days ago

a9dccc06

Clarify `*Norm` layer docstrings: `axis_index_groups` is unused under SPMD jit.

arnoegw committed 253 days ago

ae1594b0

Remove type annotation

samanklesaria committed 254 days ago

8b7c8bd5

Merge pull request #4983 from google:propagate-eager-sharding

a-googler committed 254 days ago

34e1d68e

Merge pull request #4968 from google:get-set-metadata

a-googler committed 254 days ago

b9cc05f4

Merge pull request #4981 from google:jit-wrapped-types

a-googler committed 254 days ago

609a0aba

Cristian Garcia committed 254 days ago

f8c5943c

propagate eager sharding

cgarciae committed 254 days ago

7959380f

Support shared structure

samanklesaria committed 254 days ago

126b2139

Merge pull request #4970 from google:fix-typos-flax-gspmd

a-googler committed 254 days ago

7e8f1522

Merge pull request #4982 from google:version-0.12.0

a-googler committed 254 days ago

8fe45a76

[nnx] preserve the function's type information in jit

cgarciae committed 254 days ago

566daf7d

update to version 0.12.0

cgarciae committed 254 days ago

960e52e6

Use localized YAML dumper to avoid global side effects

mohsinm-dev committed 254 days ago

Verified fee117ae

Move arrayref creation to the end of variable creation

IvyZX committed 255 days ago

6fab0a37

nnx.summary: simplify tail padding and rename test for clarity

mohsinm-dev committed 255 days ago

Verified aa5f3467

Copybara import of the project:

IvyZX committed 255 days ago

f400015e

Merge pull request #4967 from google:variable-disallow-implace-operators

a-googler committed 255 days ago

8bb57846

Add check for duplicate nodes

samanklesaria committed 256 days ago

44ed2fd9

Add compute_flops and compute_vjp options to tabulate

samanklesaria committed 256 days ago

27ed2531

add Variable.set_metadata

cgarciae committed 256 days ago

0e2b0f33

raise on Variable inplace operators

cgarciae committed 256 days ago

a4ec6a5d

Apply suggestions from code review

samanklesaria committed 256 days ago

1bd6b075

Typo fixes in flax gspmd guide

vfdev-5 committed 257 days ago

Verified e1febdad

Copybara import of the project:

epiqueras committed 260 days ago

49f51162

Merge pull request #4969 from IvyZX:cfg

a-googler committed 260 days ago

939f5323

Merge pull request #4966 from google:fix-pathparts-import

a-googler committed 260 days ago

a6ab4d7d

Newer Older