ROCm/jax

Pull Requests Commits

Wrap HIP inline functions in anonymous namespaces in vendor.h

AratiGanesh committed 66 days ago

8937a385

[ROCm] Extend GPU handle pool isolation to hipSparse and replace macros with inline functions

AratiGanesh committed 66 days ago

76b99946

Disable asan for `mgpu_torch_test`

jakeharmon8 committed 66 days ago

09dfa032

Use abstract_mesh instead of concrete_mesh when binding mesh to SDS (when pspec is passed to SDS constructor). Also refactor device_assignment extraction logic in pxla.py.

yashk2810 committed 66 days ago

44a016be

[Pallas:MGPU][NFC] Enable `test_pipelined_pipeline` for WG semantics.

allanrenucci committed 67 days ago

8f3ff0c9

Bump rules_ml_toolchain version to integrate @cuda_nvrtc header.

ybaturina committed 67 days ago

25da6a21

Merge pull request #32578 from jakevdp:abandon-tol-dep

Google-ML-Automation committed 67 days ago

0566335c

[Mosaic] Hotfix: canonicalize select to 32 bit for PF

Google-ML-Automation committed 67 days ago

16cfe37a

[JAX] Define `PyUserContextScope`

hyeontaek committed 67 days ago

ba1c18e3

pinv & matrix_rank: Abandon deprecation of old tolerance args.

jakevdp committed 67 days ago

1f930178

[pallas:sc] Avoid putting tracers into params of dma_* primitives

superbobry committed 67 days ago

c3757278

Add MPMD to IFRT IR lowerings to jaxlib._sdy_mpmd

JoelWee committed 67 days ago

2d69bd68

Merge pull request #32533 from jakevdp:dep-shard-map

Google-ML-Automation committed 67 days ago

c99f50d6

Deprecate jax.experimental.shard_map

jakevdp committed 67 days ago

d7651b64

Make `.sharding` attribute of ShapeDtypeStruct dynamic.

yashk2810 committed 67 days ago

bcead8e8

Remove use of vector.splat

Google-ML-Automation committed 67 days ago

1303b89f

Use triple overload of lookupTarget

boomanaiden154 committed 67 days ago

9f3add81

Make jax.numpy.linalg.eig return a NamedTuple.

hawkinsp committed 67 days ago

f3b4d7b5

[Pallas:MGPU] Add automatic inference for vec_size in the reduce scatter kernel

apaszke committed 67 days ago

1f3fac2a

Disable test that is failing in TPU CI.

hawkinsp committed 67 days ago

fe8daaa8

[Pallas:MGPU] Tile the minor dimensions in the reduce_scatter kernel too

apaszke committed 67 days ago

8c86b616

[Mosaic GPU][NFC] Use `Block.append` for moving operations in Mosaic GPU lowering.

allanrenucci committed 67 days ago

666e7fea

Disable failing tests in tests/layout_test.py on Cloud TPU until 2025-10-14.

hawkinsp committed 67 days ago

b15ada85

Use a single fixed channel ID for all collectives emitted by JAX.

hawkinsp committed 67 days ago

0552e7b0

[Pallas:MGPU] Add support for multicast semaphore signals

apaszke committed 67 days ago

a0a2ace0

[Pallas/interpreter] Make `shared_memory` a member of `Semaphore` (instead of asking for the global shared memory object in methods of `Semaphore`).

Google-ML-Automation committed 67 days ago

34c2fd4a

[Pallas/interpreter] Move `SharedMemory` out of the main interpreter source file.

Google-ML-Automation committed 67 days ago

9ceb3f41

[Mosaic GPU][NFC] Create undefined values in test with `ub.poison`.

allanrenucci committed 67 days ago

f44bc6e3

[mosaic] Removed usages of the deprecated `tpu.wait_dma` op

superbobry committed 67 days ago

eadee4f8

[Pallas:MGPU] Add support for tiling the transfers in reduce_scatter

apaszke committed 67 days ago

616ef4b4

Older