ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 191 days ago

4be1402b

Add test for programmatic tracing with options.

sannidhyachauhan committed 191 days ago

acf99e5e

Removing Tensorflow references from the document.

sannidhyachauhan committed 191 days ago

0f5cdba0

Add custom-call ops to roofline.

zacmustin committed 191 days ago

2dfaceda

[JAX] Relax the return type of `colocated_python` decorator

hyeontaek committed 191 days ago

b6225124

[Mosaic:TPU][NFC] Delete unused variable

tlongeri committed 191 days ago

748b39f4

Fix a missing bounds check in traceback code.

hawkinsp committed 192 days ago

cb1cc379

Merge pull request #29497 from jakevdp:make-mesh-doc

Google-ML-Automation committed 192 days ago

35ae958b

Remove legacy CPU custom calls.

dfm committed 192 days ago

8865ee66

[doc] add missing axis_types documentation

jakevdp committed 192 days ago

7aec14f1

Install SciPy from its source (head) to test against Python 3.14.0b1

kanglant committed 192 days ago

09d903fc

Set heartbeat_timeout argument and flag.

mwhittaker committed 192 days ago

a78c6a70

[export] Add back-compat test for tridiagonal solve on GPU

gnecula committed 192 days ago

e5af0880

Fix some more instances of unhashable jaxpr equation arguments.

hawkinsp committed 192 days ago

49e52c08

[jax] Increase absolute test tolerance for lax_control_flow test

basioli-k committed 192 days ago

dbde6c44

Removed unused `PyTreeDef::MakeFromNodeDataAndChildren` and its Python binding

superbobry committed 192 days ago

0e7c96a5

Bump the libtpu check to 6/20

berkinilbeyi committed 192 days ago

173574a9

Add version guards to testAutoPgle

yashk2810 committed 192 days ago

1362f7f5

Add `update_vma` and `update_weak_type` override on AbstractTMEMRef so that it return an instance of AbstractTMEMRef and not AbstractMemoryRef

yashk2810 committed 192 days ago

53849a56

PR #28102: Add cudnn paged attention support in JAX cuDNN SDPA API

Cjkkkk committed 192 days ago

639216fa

Update XLA dependency to use revision

Google-ML-Automation committed 192 days ago

b25655cf

Update XLA dependency to use revision

Google-ML-Automation committed 193 days ago

9678a764

Make mosaic_gpu equation params hashable.

hawkinsp committed 194 days ago

d065d2ac

Update XLA dependency to use revision

Google-ML-Automation committed 194 days ago

e28d6ed9

[JAX] Fix the test names in colocated_python_test.py to following the standard snake case

hyeontaek committed 194 days ago

f06888f9

Remove `with_spec` from NamedSharding and replace with `.update`

yashk2810 committed 194 days ago

7904b86b

Replace `with_partitions` and `with_unreduced` with `.update` on Partitions

yashk2810 committed 194 days ago

8e346ddb

[Pallas] Add no_pipelining debugging option to emit_pipeline.

justinjfu committed 194 days ago

dfc90529

Merge pull request #29462 from hawkinsp:nvrtc

Google-ML-Automation committed 195 days ago

a849d216

Load CUDA libraries up front with cdll.LoadLibrary().

hawkinsp committed 195 days ago

080294cf

Older