google/flax

Pull Requests Commits

This CL relaxes the input type hints to resolve an invariance issue with Python type checkers (Pyright/MyPy).

Nikitas Rontsis committed 122 days ago

287eb7b6

Merge pull request #5186 from samanklesaria:issues/5067_copy

a-googler committed 122 days ago

375c63e3

Refactor tabulate

samanklesaria committed 122 days ago

4e7147e9

chapman20j committed 123 days ago

772d9974

Merge pull request #5174 from ayulockedin:doc/audit-conv-layers

a-googler committed 123 days ago

509780a1

Merge pull request #5145 from google:fix-more-typos

a-googler committed 123 days ago

73e7985b

Merge pull request #5183 from samanklesaria/bench_fix

chapman20j committed 123 days ago

Verified 17e2b7e8

Fix bug in graph overhead benchmark

samanklesaria committed 123 days ago

514dba8d

docs(nnx): add missing functional args to Conv and LinearGeneral

ayulockedin committed 123 days ago

903fed7b

Copybara import of the project:

a-googler committed 127 days ago

a7929b37

Merge pull request #5089 from samanklesaria:sharding_metadata

a-googler committed 127 days ago

2863e252

Merge pull request #5133 from samanklesaria:issues/5130

a-googler committed 128 days ago

c4e6bca0

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

a-googler committed 128 days ago

f42707ab

Merge pull request #5115 from google:hijax-guide

a-googler committed 129 days ago

e8fe693f

improve hijax guide

cgarciae committed 129 days ago

5e31fe9e

Support multiple None and UNCONSTRAINED when resolving logical rules

marksandler2 committed 129 days ago

25610a44

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

danielsuo committed 130 days ago

5d908519

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

ayulockedin committed 130 days ago

15b2f0d7

Merge pull request #5159 from samanklesaria/pr/5140

chapman20j committed 130 days ago

Verified 4cc1d1bb

No public description

oprypin committed 130 days ago

dddd218e

fix kw_only_dataclasses for python 3.14 (part 2)

Cristian Garcia committed 130 days ago

dd15749f

samanklesaria committed 130 days ago

ebd78685

flatten respect nnx.data

cgarciae committed 130 days ago

750527aa

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

a-googler committed 130 days ago

c5398d67

Merge pull request #5160 from samanklesaria:issues/5158

a-googler committed 130 days ago

f3506885

Rename sharding_names to sharding_metadata

samanklesaria committed 130 days ago

3e963438

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

a-googler committed 130 days ago

72d5fc57

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

a-googler committed 130 days ago

a541ccb7

fix: update expected shape in conv sharding test

jackopenn committed 130 days ago

53fbca51

Add out_sharding arguments to linear layers where supported

jackopenn committed 130 days ago

23a1c1b1

Older