google/flax

Pull Requests Commits

Merge pull request #5156 from jackopenn:conv-out-shardings

a-googler committed 63 days ago

c5398d67

Merge pull request #5160 from samanklesaria:issues/5158

a-googler committed 63 days ago

f3506885

Rename sharding_names to sharding_metadata

samanklesaria committed 63 days ago

3e963438

Merge pull request #5139 from ivrolan:docs_nnx-guides-page

a-googler committed 63 days ago

72d5fc57

Merge pull request #5165 from google:fix-main-5

a-googler committed 63 days ago

a541ccb7

fix: update expected shape in conv sharding test

jackopenn committed 63 days ago

53fbca51

Add out_sharding arguments to linear layers where supported

jackopenn committed 63 days ago

23a1c1b1

ignore optax deprecation warning

cgarciae committed 63 days ago

7cca4073

Typo fix in _cached_partial method

vfdev-5 committed 63 days ago

c72bbc47

Fixed typos in the docstrings using antigravity

vfdev-5 committed 63 days ago

b006c4cd

Merge pull request #5163 from google:fix-CI-numpy-version-lt-240

a-googler committed 63 days ago

9048b4e1

Set numpy<2.4 to fix #5162

vfdev-5 committed 63 days ago

0cbe1923

Remove nnx.split_rngs call wrapping nnx.scan

samanklesaria committed 70 days ago

1aca395b

[pmap] Update `jax_utils.unreplicate` for performant use when `jax.config.jax_pmap_shmap_merge` is `True`.

danielsuo committed 75 days ago

3f2ff40e

Merge pull request #5149 from IvyZX:release

a-googler committed 80 days ago

d1791844

IvyZX committed 80 days ago

6d36623d

Merge pull request #5144 from google:fix-ci-5143

a-googler committed 81 days ago

9213cfdb

Temporary fix for failing CI

vfdev-5 committed 81 days ago

de853173

Handle pure bodies in nnx.fori_loop

samanklesaria committed 82 days ago

7255c612

fix general guides landing page

ivrolan committed 85 days ago

5a0fe50c

Merge pull request #5102 from samanklesaria:output_sharding

a-googler committed 87 days ago

ef78d658

Merge pull request #5093 from ibbyml:nnx-sdpa-parity

a-googler committed 87 days ago

be1db789

samanklesaria committed 88 days ago

66544200

Copybara import of the project:

a-googler committed 89 days ago

3139e969

Merge pull request #5057 from jorisSchaller:feat/nnx_pooling

a-googler committed 89 days ago

de84eb15

remove fingerprint

Cristian Garcia committed 90 days ago

b8aa99d4

Make popping after sowing actually remove the key

samanklesaria committed 90 days ago

a90909c9

remove abstracted_axes from nnx.jit

Cristian Garcia committed 90 days ago

b3e9b449

[jax:benchmarks] Add tracing/lowering benchmarks for a few flax examples.

danielsuo committed 90 days ago

2204c82d

[flax:examples:nlp_seq] Create a main.py file to run tests with config files to match other examples. #jax-fixit

danielsuo committed 94 days ago

872d50b8

Newer Older