google/flax

Pull Requests Commits

propagate eager sharding

cgarciae committed 96 days ago

7959380f

Support shared structure

samanklesaria committed 96 days ago

126b2139

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

a-googler committed 96 days ago

7e8f1522

Merge pull request #4982 from google:version-0.12.0

a-googler committed 96 days ago

8fe45a76

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

cgarciae committed 96 days ago

566daf7d

update to version 0.12.0

cgarciae committed 96 days ago

960e52e6

Use localized YAML dumper to avoid global side effects

mohsinm-dev committed 96 days ago

Verified fee117ae

Move arrayref creation to the end of variable creation

IvyZX committed 97 days ago

6fab0a37

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

mohsinm-dev committed 97 days ago

Verified aa5f3467

Copybara import of the project:

IvyZX committed 97 days ago

f400015e

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

a-googler committed 97 days ago

8bb57846

Add check for duplicate nodes

samanklesaria committed 98 days ago

44ed2fd9

Add compute_flops and compute_vjp options to tabulate

samanklesaria committed 98 days ago

27ed2531

add Variable.set_metadata

cgarciae committed 98 days ago

0e2b0f33

raise on Variable inplace operators

cgarciae committed 98 days ago

a4ec6a5d

Apply suggestions from code review

samanklesaria committed 98 days ago

1bd6b075

Typo fixes in flax gspmd guide

vfdev-5 committed 99 days ago

Verified e1febdad

Copybara import of the project:

epiqueras committed 102 days ago

49f51162

Merge pull request #4969 from IvyZX:cfg

a-googler committed 102 days ago

939f5323

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

a-googler committed 102 days ago

a6ab4d7d

Correctly expose flax.config.temp_flip_flag

IvyZX committed 102 days ago

ff852481

import PathParts from flax.typing

cgarciae committed 102 days ago

8485ff1e

fix Variable pytree flatten

Cristian Garcia committed 102 days ago

ccdb346a

Move iter_children and iter_modules to functions

samanklesaria committed 102 days ago

4dfafebf

Flax RNG guide for jax.jit: clarify rng outputs are shared but not inputs.

arnoegw committed 102 days ago

934bfe8c

fix HashableMapping hash definition for mixed key types

Cristian Garcia committed 103 days ago

5553f3f6

Merge pull request #4937 from google:update-bridge

a-googler committed 103 days ago

b62f1547

Merge pull request #4929 from google:pytree-guide-2

a-googler committed 103 days ago

4e8a6a72

Merge pull request #4830 from google:fix-example-gemma2

a-googler committed 103 days ago

6a75a1d3

improve pytree guide

cgarciae committed 103 days ago

bd00c17d

Newer Older