ROCm/jax

Pull Requests Commits

[ROCm] Enable previously skipped pallas tests on gfx950

amd-jianli12 committed 18 days ago

fd31e390

Use `typematch` in swap_abstract_eval instead of a shape+dtype check (it doesn't cover vma, sharding, memory_space checks)

yashk2810 committed 37 days ago

d8fc7f20

Remove the long-standing hack of trimming spec when `arr.ndim < len(spec)` (only `None`s though). Now we raise an error only in Explicit mode (Auto and Manual mode don't matter since the spec is always all `None`s).

yashk2810 committed 37 days ago

653e1d33

Update XLA dependency to use revision

kanglant committed 37 days ago

ccb2e8ad

Merge pull request #39802 from jakevdp:fix-ogrid

Google-ML-Automation committed 37 days ago

5ebb461d

[bug] fix jnp.ogrid return value

jakevdp committed 37 days ago

a7abc127

Add back call_p and closed_call_p prims to jex

yashk2810 committed 37 days ago

536b1ed3

[Pallas:SC] Add a regression test to make sure int4 kernels compile

d0k committed 38 days ago

d5784e58

Merge pull request #39593 from mattjj:push-vroyssnmnozv

Google-ML-Automation committed 38 days ago

3930db7a

Merge pull request #39640 from vfdev-5:fix-_compile_as_torch_gpu_kernel

Google-ML-Automation committed 38 days ago

be08d4c2

Merge pull request #39783 from jakevdp:meshgrid-tuple

Google-ML-Automation committed 38 days ago

e47af53f

jnp.meshgrid: return a tuple rather than a list

jakevdp committed 38 days ago

2d6f0db2

[Mosaic GPU] Define the `p2b6` MLIR dtype.

bchetioui committed 38 days ago

89c17158

[pallas:mgpu] `emit_pipeline` now works on Ampere

superbobry committed 38 days ago

ecb82891

Reverts aec6fc5ec853bb5ece7a93442fa89d18d45ca521

superbobry committed 38 days ago

2fcff0cd

[pallas:mgpu] Added a lowering rule for `mpmd_map`

superbobry committed 38 days ago

c7ed505f

Implement AOT compilation for non-autotuned JAX-Triton kernels.

EusebioDM committed 38 days ago

d35d53a7

Use new type parameter syntax in `FrozenDict`

superbobry committed 38 days ago

ca4c5a31

Remove unused layout.h/.cc from mosaic dialect.

WillFroom committed 38 days ago

3816689d

[Pallas/Mosaic GPU] Hoist `mgpu.warpgroup_barrier` out of the `if not is_cp_async` branch in `_copy_gmem_to_smem_lowering`.

bchetioui committed 38 days ago

58cce458

[mgpu] Allow `barrier_test` in warpgroup context.

chr1sj0nes committed 38 days ago

5644527b

[Mosaic TPU] Complete bf16 argmax/argmin support (on major dims and the sublane dim).

yueshengys committed 38 days ago

e2898a23

Add function name and source info to custom_vjp bwd error message when custom_vjp3 is on

yashk2810 committed 38 days ago

6cc82742

Merge pull request #39761 from Intel-tensorflow:fix-oneapi-pjrt-build-deps

Google-ML-Automation committed 38 days ago

d7774fe0

Fix wrapped_fun interacts with update_wrapper. Lead to failures because of jaxtyping

yashk2810 committed 38 days ago

07ded9f8

[oneAPI] Fix jax_oneapi_pjrt build failure due to missing @triton/intel deps

ashahba committed 38 days ago

b84f1c27

[Mosaic] Update tpu.scan_count description for ND inputs

naummo committed 38 days ago

a440810c

[Take 2] Add function name and source info to custom_vjp bwd error message

yashk2810 committed 38 days ago

85200a90

Make sure `perm` passed to ppermute is a list/tuple.

yashk2810 committed 38 days ago

d2b97cab

Access spec.partitions directly in sharding implementations.

Google-ML-Automation committed 38 days ago

053d8148

Older