ROCm/jax

Pull Requests Commits

[Mosaic GPU][NFC] Add a `is_wg_semantics` method to `mosaic_gpu_test`.

dimitar-asenov committed 39 days ago

d36ae3ba

Removed a few `type: ignore`s which were only necessary for mypy

superbobry committed 39 days ago

41f6727c

Merge pull request #35591 from superbobry:pyrefly

Google-ML-Automation committed 39 days ago

7ad75aca

Fix TensorStore implementation in Jax to use typed initialization.

Google-ML-Automation committed 39 days ago

1184dbaa

Removed unnecessary Pyrefly suppressions

superbobry committed 39 days ago

c307523f

Update XLA dependency to use revision http://github.com/openxla/xla/commit/b0037055f03bf51d364cba09e94278963aec5bcf

Google-ML-Automation committed 39 days ago

e05dc3c4

Merge pull request #35587 from gnecula:poly_bug1

Google-ML-Automation committed 40 days ago

76cbef0e

[export] Add a test for shape polymorphism with invalid constraints

gnecula committed 40 days ago

a5ed9d25

Merge pull request #35576 from jakevdp:pyrefly-opt-einsum

Google-ML-Automation committed 40 days ago

42d97c87

Fix shard_map_partial_eval and partial_eval_custom to use `nospec` instead of creating raw Pspecs

yashk2810 committed 40 days ago

165e54c7

Couple of changes in this PR:

yashk2810 committed 40 days ago

718ac012

Migrate to SafeStatic instead of SafeStaticInit.

hawkinsp committed 40 days ago

c0b26870

Merge pull request #35575 from jakevdp:pyrefly-docs

Google-ML-Automation committed 40 days ago

79d8807e

[pyrefly] test with opt-einsum

jakevdp committed 40 days ago

0711c080

[doc] update developer docs for pyrefly

jakevdp committed 40 days ago

a1328dfb

Re-apply pallas dot_general check for disallowed unsigned integer inputs.

rdyro committed 40 days ago

844bdee4

Merge pull request #35569 from superbobry:maint

Google-ML-Automation committed 40 days ago

0ef0cfa5

[Mosaic TPU] Support reshape which unfolds the minormost dim into two dims when the minormost dim is not divisible by 128, also speed up the reshape for already supported cases where the minormost dim is a multiple of 128.

yueshengys committed 40 days ago

8a73c61b

Integrate TorchTPU SDPA with optimized flash attention kernel.

jcc-google committed 40 days ago

d777b52f

Use Pyrefly instead of mypy

superbobry committed 40 days ago

b275aec2

Use driver linking as default provider.

ermilovmaxim committed 40 days ago

031c1e91

PR #35567: Bumped Pyrefly to 0.55

superbobry committed 40 days ago

3ec221c2

Merge pull request #35446 from mattjj:hi-davis

Google-ML-Automation committed 40 days ago

ecc04a88

Merge pull request #35558 from superbobry:pyrefly

Google-ML-Automation committed 40 days ago

79bdd4ba

Add support for multiple weak keys to WeakrefLRUCache.

hawkinsp committed 40 days ago

f08bb999

Rely on IFRT types implementing AbslStringify rather than calling the deprecated DebugString method

ICGog committed 40 days ago

b0efa207

Fix a pallas dot test on TPU using unsupported unsigned int dot.

rdyro committed 40 days ago

158c20ab

[hijax] enable VJPHiPrimitive to define transpose rules

mattjj committed 40 days ago

f45c878f

[pallas:mgpu] Use two barriers for try-cancel barriers in `dynamic_scheduling_loop`.

chr1sj0nes committed 40 days ago

1d8fba44

Reverts c5bf3aea3148c25e753f20b512d2efb5a741ffe9

Google-ML-Automation committed 40 days ago

5c3035c0

Newer Older