Test numba slice boxing and fix representation of None stop with negative step

Author: ricardoV94

pytensor/link/numba/dispatch/subtensor.py

@@ -46,33 +46,43 @@ def box_slice(typ, val, c):
         """
         start = c.builder.extract_value(val, 0)
         stop = c.builder.extract_value(val, 1)
+        step = c.builder.extract_value(val, 2) if typ.has_step else None
 
+        # Numba uses sys.maxsize and -sys.maxsize-1 to represent None
+        # We want to use None in the Python representation
         none_val = ir.Constant(ir.IntType(64), sys.maxsize)
+        neg_none_val = ir.Constant(ir.IntType(64), -sys.maxsize - 1)
+        none_obj = c.pyapi.get_null_object()
 
-        start_is_none = c.builder.icmp_signed("==", start, none_val)
         start = c.builder.select(
-            start_is_none,
-            c.pyapi.get_null_object(),
+            c.builder.icmp_signed("==", start, none_val),
+            none_obj,
             c.box(types.int64, start),
         )
 
-        stop_is_none = c.builder.icmp_signed("==", stop, none_val)
+        # None stop is represented as neg_none_val when step is negative
+        if step is not None:
+            stop_none_val = c.builder.select(
+                c.builder.icmp_signed(">", step, ir.Constant(ir.IntType(64), 0)),
+                none_val,
+                neg_none_val,
+            )
+        else:
+            stop_none_val = none_val
         stop = c.builder.select(
-            stop_is_none,
-            c.pyapi.get_null_object(),
+            c.builder.icmp_signed("==", stop, stop_none_val),
+            none_obj,
             c.box(types.int64, stop),
         )
 
-        if typ.has_step:
-            step = c.builder.extract_value(val, 2)
-            step_is_none = c.builder.icmp_signed("==", step, none_val)
+        if step is not None:
             step = c.builder.select(
-                step_is_none,
-                c.pyapi.get_null_object(),
+                c.builder.icmp_signed("==", step, none_val),
+                none_obj,
                 c.box(types.int64, step),
             )
         else:
-            step = c.pyapi.get_null_object()
+            step = none_obj
 
         slice_val = slice_new(c.pyapi, start, stop, step)
 

tests/link/numba/test_subtensor.py

@@ -3,7 +3,9 @@
 import numpy as np
 import pytest
 
+import pytensor.scalar as ps
 import pytensor.tensor as pt
+from pytensor import Mode, as_symbolic
 from pytensor.tensor import as_tensor
 from pytensor.tensor.subtensor import (
     AdvancedIncSubtensor,
@@ -24,6 +26,45 @@
 rng = np.random.default_rng(sum(map(ord, "Numba subtensors")))
 
 
+@pytest.mark.parametrize("step", [None, 1, 2, -2, "x"], ids=lambda x: f"step={x}")
+@pytest.mark.parametrize("stop", [None, 10, "x"], ids=lambda x: f"stop={x}")
+@pytest.mark.parametrize("start", [None, 0, 3, "x"], ids=lambda x: f"start={x}")
+def test_slice(start, stop, step):
+    x = ps.int64("x")
+
+    sym_slice = as_symbolic(
+        slice(
+            x if start == "x" else start,
+            x if stop == "x" else stop,
+            x if step == "x" else step,
+        )
+    )
+
+    no_opt_mode = Mode(linker="numba", optimizer=None)
+    evaled_slice = sym_slice.eval({x: -5}, on_unused_input="ignore", mode=no_opt_mode)
+    assert isinstance(evaled_slice, slice)
+    if start == "x":
+        assert evaled_slice.start == -5
+    elif start is None and (evaled_slice.step is None or evaled_slice.step > 0):
+        # Numba can convert to 0 (and sometimes does) in this case
+        assert evaled_slice.start in (None, 0)
+    else:
+        assert evaled_slice.start == start
+
+    if stop == "x":
+        assert evaled_slice.stop == -5
+    else:
+        assert evaled_slice.stop == stop
+
+    if step == "x":
+        assert evaled_slice.step == -5
+    elif step is None:
+        # Numba can convert to 1 (and sometimes does) in this case
+        assert evaled_slice.step in (None, 1)
+    else:
+        assert evaled_slice.step == step
+
+
 @pytest.mark.parametrize(
     "x, indices",
     [