ROCm/jax

Pull Requests Commits

adding abort logic to rocm/jax (#590)

gulsumgudukbay committed 10 days ago

a83cc9e9

Let the unit tests use build.py for setting up Bazel commands for unit tests (#582)

charleshofer committed 10 days ago

6bd1a15e

Add rocm test requirements file (#570)

AratiGanesh committed 10 days ago

835856e0

Add skip of test_is_finite() on Cuda (#565)

Arech8 committed 10 days ago

6ab35026

Address LLVM assertion failure due to a multithreaded use. Update .gitignore (#563)

Arech8 committed 10 days ago

8b0b3d10

Fix v_pages quantization and adjust test params for ROCm compatibilit… (#560)

phambinhfin committed 10 days ago

61d28322

Make nvidia version data optional for ROCm builds

phambinhfin committed 10 days ago

d88ca192

Fix hipSparse CSR algorithm mappings for ROCm 7

phambinhfin committed 10 days ago

81cdeb49

Add shared utility function get_rocm_version to test_util.py

charleshofer committed 10 days ago

1bac5e70

SWDEV-512768 - Replace hipGetLastError with hipExtGetLastError

dsicarov-amd committed 10 days ago

224b5baf

Wrap HIP inline functions in anonymous namespaces in vendor.h

AratiGanesh committed 10 days ago

2b08a20d

charleshofer committed 10 days ago

4f3741b3

hipblas typedef fix

charleshofer committed 10 days ago

6e0eb3ec

Make jaxlib targets visible

charleshofer committed 10 days ago

de91c59a

Remove nvidia_wheel_versions

charleshofer committed 10 days ago

c706c2fa

Use maxsize=None with trace_to_jaxpr's weakref_lru_cache to get more cache hits. This also restores previous behavior of lu.cache.

yashk2810 committed 13 days ago

2de5b8b0

Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.

18praveenb committed 17 days ago

98d6b4ba

Disable `tpu_splash_attention_kernel_test` on TPU v7x.

mwhittaker committed 17 days ago

28799d51

Add libtpu guard to failing tpu_trace_value_test.

mwhittaker committed 17 days ago

80b1ef6a

Skip tpu_pallas_distributed_test on 7x.

mwhittaker committed 17 days ago

77d9ffb0

Add libtpu date guard for failing test.

mwhittaker committed 17 days ago

9e7b005e

Re-enable `socket_transfer_test` internally.

emilyfertig committed 17 days ago

18b34a12

Disable `tests/multiprocess:socket_transfer_test` since it's failing flakily in all backends.

emilyfertig committed 17 days ago

d97eed6e

Fix wheel sources tests for Windows platform.

ybaturina committed 17 days ago

fae5defd

Increase shard count for TPU and GPU tests back to 5 for api_test.py.

belitskiy committed 17 days ago

8f4389a5

Prepare for JAX release 0.9.0

mwhittaker committed 18 days ago

dd745d79

Bump XLA version.

mwhittaker committed 18 days ago

cdfe904a

Add `testonly=True` to py_import targets that depend on testonly `wheel_additives` target.

ybaturina committed 18 days ago

a84ca0c5

[Pallas] Create trace_value primitive for dynamic value logging

Google-ML-Automation committed 18 days ago

02939fef

Merge pull request #34418 from jakevdp:hijax-shard-map

Google-ML-Automation committed 18 days ago

6a9f3d98

Newer Older