From 332435e0284a630a48d4a106290370275eca9408 Mon Sep 17 00:00:00 2001
From: Ruturaj4 <ruturaj.vaidya@amd.com>
Date: Tue, 2 Jul 2024 10:50:11 -0500
Subject: [PATCH] [ROCM] make mosaic dependency cuda specific

---
 jaxlib/tools/BUILD.bazel | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jaxlib/tools/BUILD.bazel b/jaxlib/tools/BUILD.bazel
index 089cba21d..8463cba08 100644
--- a/jaxlib/tools/BUILD.bazel
+++ b/jaxlib/tools/BUILD.bazel
@@ -69,11 +69,11 @@ cc_binary(
     ],
     linkshared = True,
     deps = [
-        "//jaxlib/mosaic/gpu:custom_call",
         "@xla//xla/pjrt/c:pjrt_c_api_gpu",
         "@xla//xla/pjrt/c:pjrt_c_api_gpu_version_script.lds",
         "@xla//xla/service:gpu_plugin",
     ] + if_cuda([
+        "//jaxlib/mosaic/gpu:custom_call",
         "@xla//xla/stream_executor:cuda_platform",
     ]) + if_rocm([
         "@xla//xla/stream_executor:rocm_platform",