ROCm/jax

Pull Requests Commits

Remove obsolete DCHECK on num_computations_, now that executables may have no local devices.

emilyfertig committed 1 year ago

766e68c4

Move attrs.py code to `_src/attrs.py` so that we can remove experimental imports in pjit.py which are in the middle of tracing code and can affect performance.

yashk2810 committed 1 year ago

3193ed39

PR #28157: Warn on excessive captured constants

ZacCranko committed 1 year ago

d9aee152

[docs] Remove tensorflow from tensorboard installation instructions

rdyro committed 1 year ago

8f17a552

Support `np.asarray` for JAX arrays with fully replicated shardings when some hosts have no local shards. Only hosts with local shards can fetch np arrays.

emilyfertig committed 1 year ago

6d1b5271

Replace C++ safe_zip with Python implementation using 3.10 strict keyword

Jake VanderPlas committed 1 year ago

cf976283

Add workflow for testing JAX against NumPy nightly builds.

nitins17 committed 1 year ago

13ca7002

[JAX] In McJAX, support JIT compilation/execution on only the devices attached to a single host, with other hosts not participating.

emilyfertig committed 1 year ago

268c6842

create a mosaic_gpu_support cc_library and include it as part of the //third_party/py/jax/jaxlib/cuda:cuda_gpu_kernels

Google-ML-Automation committed 1 year ago

a2d67a6a

Merge pull request #28513 from olupton:skip-cublas-12.9

Google-ML-Automation committed 1 year ago

52fcf84a

Initial f32, bf16 TPU SublanesShuffleOp

Google-ML-Automation committed 1 year ago

06988e61

Call block_until_ready for test_assert

Google-ML-Automation committed 1 year ago

5a5fd869

Call block_until_ready for testQrInvalidDtypeCPU

Google-ML-Automation committed 1 year ago

6c7e0220

jax-cuda12-plugin: require nvidia-cublas-cu12<12.9

olupton committed 1 year ago

108b288b

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

842170be

[Mosaic GPU] Enable `pallas/ops_test:test_comparison_scalar` on Mosaic GPU (except for bool)

dimitar-asenov committed 1 year ago

c75c3fae

[Mosaic GPU] Use warpgroup semantics for the FlashAttention 3 pipelined kernel.

dimitar-asenov committed 1 year ago

405c79b0

[Mosaic GPU] Run a standard mlir cse pass before Mosaic layout inference.

dimitar-asenov committed 1 year ago

63003bb5

[Mosaic GPU] A bag of fixes

apaszke committed 1 year ago

b4d8b78e

[Mosaic GPU] Improve checks in utils.ptr_to_memref

apaszke committed 1 year ago

30f38eaa

Increase tolerance in mgpu_attention_test .

jburnim committed 1 year ago

820e3dbd

Merge pull request #28461 from andportnoy:patch-4

Google-ML-Automation committed 1 year ago

d336b1b7

[XLA:Python] [JAX] Move HloPass bindings to XLA/Python.

hawkinsp committed 1 year ago

e1ca6aba

Block yt_kernel in test_upcast_to_wgmma before the assert

Google-ML-Automation committed 1 year ago

d58a7346

Add block_until_ready to buffer_callback tests where needed.

dfm committed 1 year ago

f5d621f4

Merge pull request #28500 from hawkinsp:tsan

Google-ML-Automation committed 1 year ago

ecf41a74

Merge pull request #28498 from kirillbobyrev:fix-dlpack-docs-formatting

Google-ML-Automation committed 1 year ago

e66c55aa

Improve TSAN CI build.

hawkinsp committed 1 year ago

6fa05143

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

55fb40f5

Fix `dlpack` docs formatting

kirillbobyrev committed 1 year ago

Verified b5db71bd

Older