google/flax

Pull Requests Commits

optax docs test

cgarciae committed 118 days ago

213e76b8

Merge pull request #5133 from samanklesaria:issues/5130

a-googler committed 123 days ago

c4e6bca0

Merge pull request #5166 from ayulockedin:fix-docs-typos

a-googler committed 123 days ago

f42707ab

Merge pull request #5115 from google:hijax-guide

a-googler committed 124 days ago

e8fe693f

improve hijax guide

cgarciae committed 124 days ago

5e31fe9e

Support multiple None and UNCONSTRAINED when resolving logical rules

marksandler2 committed 124 days ago

25610a44

[pmap] Avoid degraded performance under the new `jax.pmap`.

danielsuo committed 125 days ago

5d908519

docs: fix typo 'paramater' -> 'parameter'

ayulockedin committed 125 days ago

15b2f0d7

Merge pull request #5159 from samanklesaria/pr/5140

chapman20j committed 125 days ago

Verified 4cc1d1bb

No public description

oprypin committed 125 days ago

dddd218e

fix kw_only_dataclasses for python 3.14 (part 2)

Cristian Garcia committed 125 days ago

dd15749f

samanklesaria committed 125 days ago

ebd78685

flatten respect nnx.data

cgarciae committed 125 days ago

750527aa

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

a-googler committed 125 days ago

c5398d67

Merge pull request #5160 from samanklesaria:issues/5158

a-googler committed 125 days ago

f3506885

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

a-googler committed 125 days ago

72d5fc57

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

a-googler committed 125 days ago

a541ccb7

fix: update expected shape in conv sharding test

jackopenn committed 125 days ago

53fbca51

Add out_sharding arguments to linear layers where supported

jackopenn committed 125 days ago

23a1c1b1

ignore optax deprecation warning

cgarciae committed 125 days ago

7cca4073

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

a-googler committed 125 days ago

9048b4e1

Set numpy<2.4 to fix #5162

vfdev-5 committed 125 days ago

0cbe1923

Remove nnx.split_rngs call wrapping nnx.scan

samanklesaria committed 132 days ago

1aca395b

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

danielsuo committed 137 days ago

3f2ff40e

Merge pull request #5149 from IvyZX:release

a-googler committed 143 days ago

d1791844

IvyZX committed 143 days ago

6d36623d

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

a-googler committed 143 days ago

9213cfdb

Temporary fix for failing CI

vfdev-5 committed 143 days ago

de853173

fix general guides landing page

ivrolan committed 147 days ago

5a0fe50c

Merge pull request #5102 from samanklesaria:output_sharding

a-googler committed 150 days ago

ef78d658

Older