ROCm/jax

Pull Requests Commits

Introduce shim script for rocm jax upstream CI

alekstheod committed 9 days ago

009bb74a

[Pallas] Fix pipeline test to use out_type

sharadmv committed 14 days ago

a86d6336

[Pallas TPU] Add use_dummy_device_index to dma_wait_p for inserting dummy device ids in the lowering.

sharadmv committed 14 days ago

0dcd388e

replace fusible_p with Fusible(VJPHiPrimitive)

Google-ML-Automation committed 14 days ago

4a43a075

Disallow `P.UNCONSTRAINED` in `out_shardings` in `jit`. This leads to simplifications in how we lower shardings to stableHLO.

yashk2810 committed 14 days ago

b711d91b

[Pallas] Refactor pl.kernel to use out_type and scratch_types

sharadmv committed 14 days ago

71c25556

Sharding-in-types support for ragged_dot_general

rdyro committed 14 days ago

82a34e68

Merge pull request #35849 from Intel-tensorflow:oneapi_pjrt_plugin

Google-ML-Automation committed 14 days ago

9ada67a5

Sync copy single-buffered outputs in emit pipeline

Google-ML-Automation committed 14 days ago

f23ce1b5

replace fusible_p with Fusible(VJPHiPrimitive)

Cristian Garcia committed 14 days ago

e7884f87

Reduce sizes in ragged dot general vmap test

pravnar committed 14 days ago

16e3947d

Fix discharge bugs on DMA semaphores

IvyZX committed 14 days ago

1a12d2be

limit fast path to addressible single-device

Google-ML-Automation committed 14 days ago

38c31ec1

[Mosaic TPU] Support BF16 ceil, floor, sqrt, rsqrt on TPUs before v6.

yueshengys committed 14 days ago

72107c5b

[Pallas] Move CoreMemorySpace to Pallas core

sharadmv committed 14 days ago

fa737e58

Merge pull request #36613 from ymlei:add-failed-pid-to-live-devices-exception

Google-ML-Automation committed 14 days ago

50515fed

Use shaped_abstractify to handle non-jax/duck types

yashk2810 committed 14 days ago

8bb19772

Merge pull request #36984 from ROCm:rocm-skip-trisolve-perturb

Google-ML-Automation committed 14 days ago

a8bf2c11

Reverts cdcfe9d831c8b549f654984b6042c27c271a5987

Google-ML-Automation committed 14 days ago

f5eae6b7

Remove duplicate checkout in H100 CUDA workflow.

Google-ML-Automation committed 14 days ago

5c7e0eca

Fix template injection security issues in upload-test-artifacts action identified by zizmor.

Google-ML-Automation committed 14 days ago

fae53cce

Remove shape check from test

yashk2810 committed 14 days ago

e0b21fc1

Ignore more deprecation warnings in multiproc pjit test.

mwhittaker committed 14 days ago

02bb06d5

[Pallas:MGPU] Use random numbers in test.

allanrenucci committed 14 days ago

01e6ab14

Make `x: f32[4]{R:x} * scalar` work by inserting `jax.reshard` on the scalar so that it is reduced like the other operand.

yashk2810 committed 14 days ago

05b25a98

Fix security issues in download-jax-cuda-wheels action identified by zizmor.

Google-ML-Automation committed 14 days ago

414730f9

Introducing OneAPI PJRT plugin.

kranipa committed 14 days ago

5e1c8e72

Merge pull request #37028 from superbobry:maint

Google-ML-Automation committed 14 days ago

1f6a4349

Fix security issues in download-jax-rocm-wheels action identified by zizmor.

Google-ML-Automation committed 14 days ago

26d056f2

[Pallas:MGPU] Add support for slicing WGMMA accumulator under WG semantics.

allanrenucci committed 14 days ago

c9de52fa

Older