Add safe casting in Cython bindings

python/cuml/cuml/cluster/agglomerative.pyx

@@ -8,6 +8,10 @@ from cuml.common.array_descriptor import CumlArrayDescriptor
 from cuml.common.doc_utils import generate_docstring
 from cuml.internals.array import CumlArray
 from cuml.internals.base import Base, get_handle
+from cuml.internals.dimension_limits import (
+    dims_within_int_limits,
+    dims_within_size_t_limits,
+)
 from cuml.internals.mixins import ClusterMixin, CMajorInputTagMixin
 from cuml.internals.outputs import reflect
 from cuml.internals.validation import check_inputs
@@ -153,6 +157,8 @@ class AgglomerativeClustering(Base, ClusterMixin, CMajorInputTagMixin):
             ensure_min_samples=2,
             reset=True,
         )
+        dims_within_int_limits(n_rows=X.shape[0], n_cols=X.shape[1])
+        dims_within_size_t_limits(n_clusters=self.n_clusters)
         cdef int n_rows = X.shape[0]
         cdef int n_cols = X.shape[1]
 

python/cuml/cuml/cluster/dbscan.pyx

@@ -9,6 +9,10 @@ from cuml.common.doc_utils import generate_docstring
 from cuml.internals import logger, reflect
 from cuml.internals.array import CumlArray
 from cuml.internals.base import Base, get_handle
+from cuml.internals.dimension_limits import (
+    dims_within_int_limits,
+    dims_within_size_t_limits,
+)
 from cuml.internals.interop import (
     InteropMixin,
     UnsupportedOnGPU,
@@ -336,6 +340,9 @@ class DBSCAN(Base,
         cdef int64_t n_rows = X.shape[0]
         cdef int64_t n_cols = X.shape[1]
 
+        dims_within_size_t_limits(n_rows=n_rows, n_cols=n_cols)
+        dims_within_int_limits(min_samples=self.min_samples)
+
         if out_dtype not in (cp.dtype("int32"), cp.dtype("int64")):
             raise ValueError(
                 f"Expected out_dtype to be one of ['int32', 'int64'], got {out_dtype!s}"

python/cuml/cuml/cluster/hdbscan/hdbscan.pyx

@@ -9,6 +9,10 @@ from cuml.common.doc_utils import generate_docstring
 from cuml.internals import logger, reflect
 from cuml.internals.array import CumlArray
 from cuml.internals.base import Base, get_handle
+from cuml.internals.dimension_limits import (
+    dims_within_int_limits,
+    dims_within_size_t_limits,
+)
 from cuml.internals.interop import (
     InteropMixin,
     UnsupportedOnGPU,
@@ -145,6 +149,8 @@ cdef class _HDBSCANState:
         lambdas = np.ascontiguousarray(tree["lambda_val"], dtype=np.float32)
         sizes = np.ascontiguousarray(tree["child_size"], dtype=np.int64)
 
+        dims_within_int_limits(n_edges=len(tree))
+        dims_within_size_t_limits(n_leaves=n_leaves)
         cdef int n_edges = len(tree)
         cdef handle_t *handle_ = <handle_t*> <size_t> handle.getHandle()
         self.condensed_tree = new lib.CondensedHierarchy[int64_t, float](
@@ -181,6 +187,7 @@ cdef class _HDBSCANState:
 
         cdef _HDBSCANState self = _HDBSCANState.__new__(_HDBSCANState)
 
+        dims_within_int_limits(n_rows=X.shape[0], n_cols=X.shape[1])
         cdef int n_rows = X.shape[0]
         cdef int n_cols = X.shape[1]
 
@@ -262,6 +269,8 @@ cdef class _HDBSCANState:
         sizes = cp.asarray(dendrogram[:, 3], order="C", dtype="int64")
 
         cdef size_t n_leaves = dendrogram.shape[0] + 1
+        dims_within_size_t_limits(n_leaves=n_leaves, dendrogram_rows=dendrogram.shape[0])
+        dims_within_int_limits(min_cluster_size=min_cluster_size)
 
         handle = get_handle()
         cdef handle_t *handle_ = <handle_t*> <size_t> handle.getHandle()
@@ -293,6 +302,7 @@ cdef class _HDBSCANState:
         """Initialize internal state from a new `fit`"""
         cdef _HDBSCANState self = _HDBSCANState.__new__(_HDBSCANState)
 
+        dims_within_int_limits(n_rows=X.shape[0], n_cols=X.shape[1])
         cdef int n_rows = X.shape[0]
         cdef int n_cols = X.shape[1]
 
@@ -404,6 +414,7 @@ cdef class _HDBSCANState:
 
         handle = get_handle()
 
+        dims_within_int_limits(n_rows=X.shape[0], n_cols=X.shape[1])
         cdef int n_rows = X.shape[0]
         cdef int n_cols = X.shape[1]
         cdef int64_t* labels_ptr = <int64_t*><uintptr_t>labels.ptr
@@ -1227,6 +1238,11 @@ def membership_vector(clusterer, points_to_predict, int batch_size=4096, convert
         order="C",
         return_index=True,
     )
+    dims_within_int_limits(
+        n_prediction_points=points_to_predict.shape[0],
+        n_clusters=clusterer.n_clusters_,
+        batch_size=batch_size,
+    )
     cdef int n_prediction_points = points_to_predict.shape[0]
 
     membership_vec = cp.empty(
@@ -1312,6 +1328,7 @@ def approximate_predict(clusterer, points_to_predict, convert_dtype=True):
         order="C",
         return_index=True,
     )
+    dims_within_int_limits(n_prediction_points=points_to_predict.shape[0])
     cdef int n_prediction_points = points_to_predict.shape[0]
 
     prediction_labels = cp.empty(n_prediction_points, dtype="int64")
@@ -1391,6 +1408,8 @@ def _extract_clusters(
     Exposed for testing only"""
     cdef size_t n_leaves = condensed_tree["parent"].min()
     cdef int n_edges = len(condensed_tree)
+    dims_within_int_limits(n_edges=n_edges)
+    dims_within_size_t_limits(n_leaves=n_leaves)
 
     parents = cp.asarray(condensed_tree["parent"], order="C", dtype="int64")
     children = cp.asarray(condensed_tree["child"], order="C", dtype="int64")

python/cuml/cuml/cluster/kmeans.pyx

@@ -8,6 +8,7 @@ from cuml.common.array_descriptor import CumlArrayDescriptor
 from cuml.common.doc_utils import generate_docstring
 from cuml.internals.array import CumlArray
 from cuml.internals.base import Base, get_handle
+from cuml.internals.dimension_limits import INT32_MAX, dims_within_int_limits
 from cuml.internals.interop import (
     InteropMixin,
     UnsupportedOnGPU,
@@ -35,6 +36,12 @@ cdef _kmeans_init_params(kmeans, lib.KMeansParams& params):
     """Initialize a passed KMeansParams instance from a KMeans instance."""
     cdef bool multi_gpu = kmeans._multi_gpu
 
+    dims_within_int_limits(
+        n_clusters=kmeans.n_clusters,
+        max_samples_per_batch=kmeans.max_samples_per_batch,
+        max_iter=kmeans.max_iter,
+    )
+
     params.n_clusters = kmeans.n_clusters
     params.max_iter = kmeans.max_iter
     params.tol = kmeans.tol
@@ -78,6 +85,8 @@ cdef _kmeans_init_params(kmeans, lib.KMeansParams& params):
     else:
         params.n_init = kmeans.n_init
 
+    dims_within_int_limits(n_init=params.n_init)
+
 
 cdef _kmeans_fit(
     handle_t& handle,
@@ -91,7 +100,11 @@ cdef _kmeans_fit(
     cdef int64_t n_cols = X.shape[1]
 
     cdef bool values_f32 = X.dtype == cp.float32
-    cdef bool indices_i32 = (n_rows * n_cols) < (2**31 - 1)
+    cdef bool use_i32_dims = (
+        n_rows <= INT32_MAX
+        and n_cols <= INT32_MAX
+        and n_rows * n_cols < INT32_MAX
+    )
 
     cdef uintptr_t X_ptr = X.data.ptr
     cdef uintptr_t centers_ptr = centers.data.ptr
@@ -104,7 +117,7 @@ cdef _kmeans_fit(
 
     with nogil:
         if values_f32:
-            if indices_i32:
+            if use_i32_dims:
                 lib.fit(
                     handle,
                     params,
@@ -129,7 +142,7 @@ cdef _kmeans_fit(
                     n_iter_64,
                 )
         else:
-            if indices_i32:
+            if use_i32_dims:
                 lib.fit(
                     handle,
                     params,
@@ -153,7 +166,7 @@ cdef _kmeans_fit(
                     inertia_64,
                     n_iter_64,
                 )
-    return n_iter_32 if indices_i32 else n_iter_64
+    return n_iter_32 if use_i32_dims else n_iter_64
 
 
 cdef _kmeans_predict(
@@ -170,9 +183,15 @@ cdef _kmeans_predict(
     cdef int64_t n_rows = X.shape[0]
     cdef int64_t n_cols = X.shape[1]
 
+    cdef bool use_i32_dims = (
+        n_rows <= INT32_MAX
+        and n_cols <= INT32_MAX
+        and n_rows * n_cols < INT32_MAX
+    )
+
     labels = cp.zeros(
         shape=n_rows,
-        dtype=(cp.int32 if n_rows * n_cols < 2**31 - 1 else cp.int64),
+        dtype=(cp.int32 if use_i32_dims else cp.int64),
     )
 
     cdef uintptr_t X_ptr = X.data.ptr
@@ -181,14 +200,13 @@ cdef _kmeans_predict(
     cdef uintptr_t labels_ptr = labels.data.ptr
 
     cdef bool values_f32 = X.dtype == cp.float32
-    cdef bool indices_i32 = labels.dtype == cp.int32
 
     cdef float inertia_f32 = 0
     cdef double inertia_f64 = 0
 
     with nogil:
         if values_f32:
-            if indices_i32:
+            if use_i32_dims:
                 lib.predict(
                     handle,
                     params,
@@ -215,7 +233,7 @@ cdef _kmeans_predict(
                     inertia_f32,
                 )
         else:
-            if indices_i32:
+            if use_i32_dims:
                 lib.predict(
                     handle,
                     params,
@@ -687,6 +705,12 @@ class KMeans(Base,
         cdef int64_t n_rows = X.shape[0]
         cdef int64_t n_cols = X.shape[1]
 
+        cdef bool use_i32_dims = (
+            n_rows <= INT32_MAX
+            and n_cols <= INT32_MAX
+            and n_rows * n_cols < INT32_MAX
+        )
+
         out = cp.zeros(
             shape=(n_rows, self.n_clusters), dtype=X.dtype, order="C",
         )
@@ -701,11 +725,10 @@ class KMeans(Base,
         _kmeans_init_params(self, params)
 
         cdef bool values_f32 = X.dtype == cp.float32
-        cdef bool indices_i32 = self.labels_.dtype == cp.int32
 
         with nogil:
             if values_f32:
-                if indices_i32:
+                if use_i32_dims:
                     lib.transform(
                         handle_[0],
                         params,
@@ -726,7 +749,7 @@ class KMeans(Base,
                         <float*>out_ptr,
                     )
             else:
-                if indices_i32:
+                if use_i32_dims:
                     lib.transform(
                         handle_[0],
                         params,

python/cuml/cuml/cluster/spectral_clustering.pyx

@@ -9,6 +9,7 @@ import cuml
 from cuml.common.array_descriptor import CumlArrayDescriptor
 from cuml.internals.array import CumlArray
 from cuml.internals.base import Base, get_handle
+from cuml.internals.dimension_limits import dims_within_int_limits
 from cuml.internals.interop import (
     InteropMixin,
     UnsupportedOnGPU,
@@ -303,8 +304,10 @@ class SpectralClustering(Base,
             ensure_min_samples=2,
             reset=True,
         )
-        cdef int n_samples, n_features
-        n_samples, n_features = X.shape
+        n_samples_py, n_features_py = map(int, X.shape)
+        dims_within_int_limits(n_samples=n_samples_py, n_features=n_features_py)
+        cdef int n_samples = n_samples_py
+        cdef int n_features = n_features_py
 
         cdef float* affinity_data_ptr = NULL
         cdef int* affinity_rows_ptr = NULL
@@ -351,6 +354,14 @@ class SpectralClustering(Base,
         config.n_components = max(1, min(effective_n_components, (n_samples - 1) // 3))
         config.n_neighbors = min(self.n_neighbors, n_samples - 1)
         config.n_init = self.n_init
+        if precomputed:
+            dims_within_int_limits(affinity_nnz=int(affinity_nnz))
+        dims_within_int_limits(
+            n_clusters=self.n_clusters,
+            n_init=self.n_init,
+            n_components=int(config.n_components),
+            n_neighbors=int(config.n_neighbors),
+        )
         if self.eigen_tol == "auto":
             config.eigen_tol = 0.0
         else:

python/cuml/cuml/common/opg_data_utils_mg.pyx

@@ -1,12 +1,16 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
 
 import numpy as np
 
 from cuml.common import input_to_cuml_array
 from cuml.internals.array import CumlArray
+from cuml.internals.dimension_limits import (
+    dims_within_int_limits,
+    dims_within_size_t_limits,
+)
 
 from cython.operator cimport dereference as deref
 from libc.stdint cimport uintptr_t
@@ -112,10 +116,12 @@ def build_rank_size_pair(parts_to_sizes, rank):
     cdef vector[RankSizePair*] *rsp_vec = new vector[RankSizePair*]()
 
     for idx, rankToSize in enumerate(parts_to_sizes):
-        rank, size = rankToSize
+        part_rank, part_size = rankToSize
+        dims_within_int_limits(rank=part_rank)
+        dims_within_size_t_limits(partition_size=part_size)
         rsp = <RankSizePair*> malloc(sizeof(RankSizePair))
-        rsp.rank = <int>rank
-        rsp.size = <size_t>size
+        rsp.rank = <int>part_rank
+        rsp.size = <size_t>part_size
 
         rsp_vec.push_back(rsp)
 
@@ -157,6 +163,9 @@ def build_part_descriptor(m, n, rank_size_t, rank):
     --------
     ptr: PartDescriptor object
     """
+    dims_within_size_t_limits(total_rows=m, n_cols=n)
+    dims_within_int_limits(rank=rank)
+
     cdef uintptr_t rank_size_ptr = rank_size_t
 
     cdef vector[RankSizePair *] *rsp_vec \
@@ -200,6 +209,8 @@ def _build_part_inputs(cuda_arr_ifaces,
                        parts_to_ranks,
                        m, n, local_rank,
                        convert_dtype):
+    dims_within_size_t_limits(total_rows=m, n_cols=n)
+    dims_within_int_limits(local_rank=local_rank)
 
     cuml_arr_ifaces = []
     for arr in cuda_arr_ifaces:
@@ -220,10 +231,12 @@ def _build_part_inputs(cuda_arr_ifaces,
 
     cdef vector[RankSizePair*] partsToRanks
     for idx, rankToSize in enumerate(parts_to_ranks):
-        rank, size = rankToSize
+        part_rank, part_size = rankToSize
+        dims_within_int_limits(rank=part_rank)
+        dims_within_size_t_limits(partition_size=part_size)
         rsp = <RankSizePair*>malloc(sizeof(RankSizePair))
-        rsp.rank = <int>rank
-        rsp.size = <size_t>size
+        rsp.rank = <int>part_rank
+        rsp.size = <size_t>part_size
         partsToRanks.push_back(rsp)
 
     cdef PartDescriptor *descriptor = \