ROCm/jax

Pull Requests Commits

Mount /home/runner/_work in ROCm workflow containers

Kunreddy committed 78 days ago

b5d9ff49

[ROCm] Manually bump of XLA pin.

mminutoli committed 78 days ago

daadb3b0

[ROCm] Add expanded target arch set for ROCm to rocm.bazelrc.

tsrw2048 committed 78 days ago

9dbd68a2

[ROCm] nightly: build jax/jaxlib on fork runners

mminutoli committed 78 days ago

8654db62

Add mosaic lowering rule for eval_jaxpr_p

yashk2810 committed 78 days ago

a6aadfa7

Merge pull request #38064 from aybchan:aybchan/reduce-array-print-out

Google-ML-Automation committed 78 days ago

2f4db679

Merge pull request #38517 from aybchan:aybchan/scaled-matmul-total-devices

Google-ML-Automation committed 78 days ago

23bd744e

Update XLA dependency to use revision

Google-ML-Automation committed 78 days ago

71742279

Get rid of DceSink as a VJPHiPrimitive because in `expand` when `prevent_mlir_dce` is False, we still need a dce_sink lojax primitive so that we don't return an empty tuple and allow JAX level DCE to get rid of dce_sink.

yashk2810 committed 78 days ago

7ad391e6

[pallas:mgpu] Automatically infer `PROMISE_IN_BOUNDS` in `emit_pipeline`

superbobry committed 78 days ago

4861a40d

Replace closed form expansion of 2x2 and 3x3 determinants with a closed form LU decomposition with row pivoting.

hawkinsp committed 78 days ago

a25d587e

[triton] Added support for host-side TMA descriptors

superbobry committed 78 days ago

8585ff73

[mosaic_gpu] The `barrier` operand to `async_load` is now optional on pre-Hopper GPUs

superbobry committed 78 days ago

263eed93

Implement AOT compilation for autotuned JAX-Triton kernels.

EusebioDM committed 79 days ago

8e07bd40

Merge pull request #39895 from jax-ml:bzlmod

Google-ML-Automation committed 79 days ago

51b5b00f

[pallas] Deprecated `pl.core_map` in favor of `pl.kernel`

superbobry committed 79 days ago

d3287b70

Switch to bzlmod by default.

hawkinsp committed 79 days ago

aeff4e4b

Merge pull request #39901 from jax-ml:reqs

Google-ML-Automation committed 79 days ago

24bd68dd

Regenerate requirements lock files.

hawkinsp committed 79 days ago

aa742407

Add emulation for int4 cmpi

SamGinzburg committed 79 days ago

798b72fa

Merge pull request #39915 from mattjj:push-ustwymvzzpon

Google-ML-Automation committed 79 days ago

da2d20de

Make dce_sink a VJPHiPrimitive so that it works with sharding

yashk2810 committed 79 days ago

4b315c97

[custom_vjp3] fix export + gspmdsharding + custom_vjp3

mattjj committed 79 days ago

2f55e320

Merge pull request #39606 from Intel-tensorflow:oneapi_callback_support

Google-ML-Automation committed 79 days ago

5a55d303

Merge pull request #39912 from mattjj:push-vrrrxxwtyyor

Google-ML-Automation committed 79 days ago

cc77b3bb

[custom_vjp3] misc fixes, in-tree tests now green

mattjj committed 79 days ago

27b42375

Merge pull request #39765 from ROCm:drop-rocm-sdk-path-block-upstream

Google-ML-Automation committed 79 days ago

d46fcd96

Fix ASan memory leak in PJRT custom call handler registration

akuegel committed 79 days ago

90eac75f

[pallas:mgpu] Use `mpmd_map` instead of `core_map` in `plgpu.warp_map`

superbobry committed 79 days ago

79616548

Merge pull request #39909 from mattjj:push-wpyxnttptsxm

Google-ML-Automation committed 79 days ago

593a0203

Older