ROCm/jax

Pull Requests Commits

Remove test scripts

alekstheod committed 3 hours ago

d3bd9c46

Remove unused import os

alekstheod committed 3 hours ago

b8099099

Simplify _get_library_path to only check for xla_rocm_plugin.so

alekstheod committed 3 hours ago

e5e7c2c5

Move ROCm GPU detection to hardware_utils and add multi-GPU shm warning

alekstheod committed 3 hours ago

74e8d739

[Mosaic] Generate enums in Python bindings

tlongeri committed 5 hours ago

ebf39ff6

jaxlib: Expose the status attributes on XlaRuntimeError.

pizzud committed 6 hours ago

ffcc5ea9

[jaxlib] Fixed usages of `typing_extensions.CapsuleType` in stubs

superbobry committed 6 hours ago

ed1e8d37

Merge pull request #35713 from alekstheod:patch-2

Google-ML-Automation committed 7 hours ago

44e12c3a

Merge pull request #35747 from jakevdp:ai-policy

Google-ML-Automation committed 7 hours ago

719aaea1

[jaxlib] Added a missing include and regenerated the .pyis

superbobry committed 7 hours ago

062845f0

Increase size of gpu_pallas_interpret_test_cpu to avoid timeout.

mwhittaker committed 8 hours ago

3853b241

Expose register_extra_dialect in mosaic to support loading custom dialect attributes for backward compatability.

Google-ML-Automation committed 8 hours ago

78126876

Increase size of `lax_scipy_spectral_dac_test` to avoid timeouts.

mwhittaker committed 8 hours ago

3f6a1731

Merge pull request #35748 from jakevdp:pyrefly-sparse

Google-ML-Automation committed 8 hours ago

da78e804

[MGPU] Add AsyncStoreSmemToTmem Op.

golechwierowicz committed 8 hours ago

336a1de7

[Mosaic GPU] Add support for reducing strided and splat reduce layout constaints.

allanrenucci committed 8 hours ago

18a19f8b

[Pallas/Mosaic GPU] Enable warpgroup semantics-based tests for the `HelpersTest` and `DistributedTest` classes.

bchetioui committed 8 hours ago

0453525b

[Pallas] Add more test cases for folding last two dims reshape.

WindQAQ committed 8 hours ago

2b93660d

[Pallas/Mosaic GPU] Enable lowering `semaphore_{read,signal,wait}_p` using warpgroup semantics.

bchetioui committed 9 hours ago

4e9175fb

Remove extra space in JAX OSS workflow job names.

mwhittaker committed 9 hours ago

cc60de73

Merge pull request #35710 from alekstheod:patch-1

Google-ML-Automation committed 10 hours ago

212235b3

[pyrefly] add type checking for jax/experimental/sparse

jakevdp committed 11 hours ago

f722be11

[Pallas:TPU] Use constant 0 indices for grid dims known to be 1

apaszke committed 11 hours ago

7da7f39c

Fix rank mismatches in Pallas memory reference transformations.

cperivol committed 11 hours ago

d1e9445b

[Pallas:MGPU] Make sure the ue8m0 reciprocal optimization kicks in

apaszke committed 12 hours ago

fd7de07b

Merge pull request #35716 from olupton:re-enable-test

Google-ML-Automation committed 12 hours ago

c9d722f0

[Pallas][TPU kernel interpreter] Make behaviour of `LoggingInfo` TPU/GPU specific.

Google-ML-Automation committed 12 hours ago

9d1c6a04

[Mosaic GPU] Add a leader_tracked parameter in async_copy

apaszke committed 13 hours ago

89b7d339

[mgpu] Fix for mbarriers with `order_tensor_core=True`.

chr1sj0nes committed 14 hours ago

acc3a5b4

[Mosaic TPU] add sign lowering

cota committed 14 hours ago

b6113e78

Older