From 180ae55b3579f7bbd888fe98b98887e7ec000913 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Sat, 6 Dec 2025 14:53:07 +0100 Subject: [PATCH] Numba RavelMultiIndex: Fix scalars with clip mode --- pytensor/link/numba/dispatch/extra_ops.py | 6 +++++- tests/link/numba/test_extra_ops.py | 12 ++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/pytensor/link/numba/dispatch/extra_ops.py b/pytensor/link/numba/dispatch/extra_ops.py index 21c311b0f3..105d2fc2d1 100644 --- a/pytensor/link/numba/dispatch/extra_ops.py +++ b/pytensor/link/numba/dispatch/extra_ops.py @@ -154,7 +154,11 @@ def ravelmultiindex(*inp): stacked_indices[..., i] %= dim_limit elif mode == "clip": dim_indices = stacked_indices[..., i] - stacked_indices[..., i] = np.clip(dim_indices, 0, dim_limit - 1) + # Cannot call np.clip on scalars + if vec_indices: + stacked_indices[..., i] = np.clip(dim_indices, 0, dim_limit - 1) + else: + stacked_indices[..., i] = max(0, min(dim_indices, dim_limit - 1)) else: # raise dim_indices = stacked_indices[..., i] invalid_indices = (dim_indices < 0) | (dim_indices >= shape[i]) diff --git a/tests/link/numba/test_extra_ops.py b/tests/link/numba/test_extra_ops.py index 5b523d5eb4..7375ce9396 100644 --- a/tests/link/numba/test_extra_ops.py +++ b/tests/link/numba/test_extra_ops.py @@ -171,6 +171,12 @@ def test_FillDiagonalOffset(a, val, offset): "raise", ValueError, ), + ( + tuple((pt.lscalar(), v) for v in np.array([0, 0, 3])), + (pt.lvector(), np.array([2, 3, 4])), + "wrap", + None, + ), ( tuple( (pt.lvector(), v) for v in np.array([[0, 1, 2], [2, 0, 3], [1, 3, 5]]) @@ -188,6 +194,12 @@ def test_FillDiagonalOffset(a, val, offset): "wrap", None, ), + ( + tuple((pt.lscalar(), v) for v in np.array([0, 0, 3])), + (pt.lvector(), np.array([2, 3, 4])), + "clip", + None, + ), ( tuple( (pt.lvector(), v) for v in np.array([[0, 1, 2], [2, 0, 3], [1, 3, 5]])