From 1c4ae9ac4d5a39d13ec32ebf2095912dce76d991 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Longeri?= <tlongeri@google.com>
Date: Tue, 21 Nov 2023 15:57:35 -0800
Subject: [PATCH] Fix bad cast in tpu_ext.cc

The argument to the cast value is of type ssize_t. Mismatch between int64_t and ssize_t happens in Mac and causes build to fail:
`error: const_cast from 'const pybind11::ssize_t *' (aka 'const long *') to 'int64_t *' (aka 'long long *') is not allowed`

PiperOrigin-RevId: 584451928
---
 jaxlib/mlir/_mlir_libs/tpu_ext.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jaxlib/mlir/_mlir_libs/tpu_ext.cc b/jaxlib/mlir/_mlir_libs/tpu_ext.cc
index 174c493c2382..efc0597e489d 100644
--- a/jaxlib/mlir/_mlir_libs/tpu_ext.cc
+++ b/jaxlib/mlir/_mlir_libs/tpu_ext.cc
@@ -576,7 +576,7 @@ PYBIND11_MODULE(_tpu_ext, m) {
           if constexpr (std::is_same_v<ssize_t, int64_t>) {
             // TODO(tlongeri): Find a way to avoid the const_cast
             return mlirTpuAssemble(getDefaultInsertionPoint(), ty, layout,
-                                   {{const_cast<int64_t*>(np_arr.shape()),
+                                   {{const_cast<ssize_t*>(np_arr.shape()),
                                      static_cast<size_t>(np_arr.ndim())},
                                     vals.data()},
                                    TARGET_SHAPE);