ROCm/jax

Pull Requests Commits

charleshofer committed 56 days ago

b57301f4

charleshofer committed 56 days ago

524a110c

Use actinolite cluster

charleshofer committed 56 days ago

ac2793c8

Update XLA dependency to use revision

Google-ML-Automation committed 56 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 56 days ago

7ad391e6

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

superbobry committed 57 days ago

4861a40d

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

hawkinsp committed 57 days ago

a25d587e

[triton] Added support for host-side TMA descriptors

superbobry committed 57 days ago

8585ff73

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

superbobry committed 57 days ago

263eed93

Implement AOT compilation for autotuned JAX-Triton kernels.

EusebioDM committed 57 days ago

8e07bd40

Merge pull request #39895 from jax-ml:bzlmod

Google-ML-Automation committed 57 days ago

51b5b00f

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

superbobry committed 57 days ago

d3287b70

Switch to bzlmod by default.

hawkinsp committed 57 days ago

aeff4e4b

Merge pull request #39901 from jax-ml:reqs

Google-ML-Automation committed 57 days ago

24bd68dd

Regenerate requirements lock files.

hawkinsp committed 57 days ago

aa742407

Add emulation for int4 cmpi

SamGinzburg committed 57 days ago

798b72fa

Merge pull request #39915 from mattjj:push-ustwymvzzpon

Google-ML-Automation committed 57 days ago

da2d20de

Make dce_sink a VJPHiPrimitive so that it works with sharding

yashk2810 committed 57 days ago

4b315c97

[custom_vjp3] fix export + gspmdsharding + custom_vjp3

mattjj committed 57 days ago

2f55e320

Merge pull request #39606 from Intel-tensorflow:oneapi_callback_support

Google-ML-Automation committed 57 days ago

5a55d303

Merge pull request #39912 from mattjj:push-vrrrxxwtyyor

Google-ML-Automation committed 57 days ago

cc77b3bb

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

mattjj committed 57 days ago

27b42375

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

Google-ML-Automation committed 57 days ago

d46fcd96

Fix ASan memory leak in PJRT custom call handler registration

akuegel committed 57 days ago

90eac75f

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

superbobry committed 57 days ago

79616548

Merge pull request #39909 from mattjj:push-wpyxnttptsxm

Google-ML-Automation committed 57 days ago

593a0203

[remat3] fix vmap axes=None bug

mattjj committed 57 days ago

a6e59252

Merge pull request #39340 from wangfakang:opt_topo_aware

Google-ML-Automation committed 57 days ago

99783a3a

Wrap negative indices by default in take_along_axis and remove legacy behavior that used promise_in_bounds.

18praveenb committed 57 days ago

bbf826af

add python and buffer callback support for oneapi

kranipa committed 57 days ago

1f3308d0

Older