Converts negative constant axis to positive if present in Join(COp)

pytensor/tensor/basic.py

@@ -2471,6 +2471,18 @@ def make_node(self, axis, *tensors):
         if axis.type.ndim > 0:
             raise TypeError(f"Axis {axis} must be 0-d.")
 
+        # Convert negative constant axis to positive during canonicalization
+        if isinstance(axis, Constant) and tensors:
+            # Get the axis value directly from the constant's data
+            axis_val = axis.data.item()
+            # Check if it's negative and needs normalization
+            if axis_val < 0:
+                ndim = tensors[0].ndim
+                # Convert negative axis to positive
+                axis_val = normalize_axis_index(axis_val, ndim)
+                # Replace the original axis with the normalized one
+                axis = constant(axis_val, dtype=axis.type.dtype)
+
         tensors = [as_tensor_variable(x) for x in tensors]
 
         if not builtins.all(targs.type.ndim > 0 for targs in tensors):

tests/tensor/test_basic.py

@@ -2179,6 +2179,43 @@ def test_join_performance(self, ndim, axis, memory_layout, gc, benchmark):
         assert fn(*test_values).shape == (n * 6, n)[:ndim] if axis == 0 else (n, n * 6)
         benchmark(fn, *test_values)
 
+    def test_join_negative_axis_rewrite(self):
+        """Test that constant negative axis is rewritten to positive axis during canonicalization."""
+        v = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], dtype=self.floatX)
+        a = self.shared(v)
+        b = as_tensor_variable(v)
+
+        # Create join with negative axis
+        s = join(-1, a, b)
+
+        # Get the actual Join op node from the graph
+        f = pytensor.function([], [s], mode=self.mode)
+
+        # Directly access the Join node from the output's owner
+        join_node = f.maker.fgraph.outputs[0].owner
+        assert isinstance(join_node.op, Join), "Expected output node to be a Join op"
+
+        # Check that the axis input has been converted to a constant with value 1 (not -1)
+        axis_input = join_node.inputs[0]
+        assert isinstance(axis_input, ptb.Constant), "Expected axis to be a Constant"
+        assert (
+            axis_input.data == 1
+        ), f"Expected axis to be normalized to 1, got {axis_input.data}"
+
+        # Now test with axis -2 which should be rewritten to 0
+        s2 = join(-2, a, b)
+        f2 = pytensor.function([], [s2], mode=self.mode)
+
+        join_node = f2.maker.fgraph.outputs[0].owner
+        assert isinstance(join_node.op, Join), "Expected output node to be a Join op"
+
+        # Check that the axis input has been converted to a constant with value 0 (not -2)
+        axis_input = join_node.inputs[0]
+        assert isinstance(axis_input, ptb.Constant), "Expected axis to be a Constant"
+        assert (
+            axis_input.data == 0
+        ), f"Expected axis to be normalized to 0, got {axis_input.data}"
+
 
 def test_TensorFromScalar():
     s = ps.constant(56)