ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4be1402b

Add test for programmatic tracing with options.

sannidhyachauhan committed 1 year ago

acf99e5e

Removing Tensorflow references from the document.

sannidhyachauhan committed 1 year ago

0f5cdba0

Add custom-call ops to roofline.

zacmustin committed 1 year ago

2dfaceda

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

hyeontaek committed 1 year ago

b6225124

[Mosaic:TPU][NFC] Delete unused variable

tlongeri committed 1 year ago

748b39f4

Fix a missing bounds check in traceback code.

hawkinsp committed 1 year ago

cb1cc379

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

Google-ML-Automation committed 1 year ago

35ae958b

Remove legacy CPU custom calls.

dfm committed 1 year ago

8865ee66

[doc] add missing axis_types documentation

jakevdp committed 1 year ago

7aec14f1

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

kanglant committed 1 year ago

09d903fc

Set heartbeat_timeout argument and flag.

mwhittaker committed 1 year ago

a78c6a70

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

gnecula committed 1 year ago

e5af0880

Fix some more instances of unhashable jaxpr equation arguments.

hawkinsp committed 1 year ago

49e52c08

[jax] Increase absolute test tolerance for lax_control_flow test

basioli-k committed 1 year ago

dbde6c44

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

superbobry committed 1 year ago

0e7c96a5

Bump the libtpu check to 6/20

berkinilbeyi committed 1 year ago

173574a9

Add version guards to testAutoPgle

yashk2810 committed 1 year 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 1 year ago

53849a56

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

Cjkkkk committed 1 year ago

639216fa

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b25655cf

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

9678a764

Make mosaic_gpu equation params hashable.

hawkinsp committed 1 year ago

d065d2ac

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

e28d6ed9

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

hyeontaek committed 1 year ago

f06888f9

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

yashk2810 committed 1 year ago

7904b86b

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

yashk2810 committed 1 year ago

8e346ddb

[Pallas] Add no_pipelining debugging option to emit_pipeline.

justinjfu committed 1 year ago

dfc90529

Merge pull request #29462 from hawkinsp:nvrtc

Google-ML-Automation committed 1 year ago

a849d216

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

hawkinsp committed 1 year ago

080294cf

Older