From a1f99bf04869e74952981bd81fab3198a547f419 Mon Sep 17 00:00:00 2001 From: PabloRoque Date: Tue, 8 Jul 2025 10:14:28 +0200 Subject: [PATCH 1/7] Vibe coding #1505 --- pytensor/tensor/basic.py | 12 ++++++++++ tests/tensor/test_basic.py | 45 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+) diff --git a/pytensor/tensor/basic.py b/pytensor/tensor/basic.py index 5a3cf0036f..aba6787826 100644 --- a/pytensor/tensor/basic.py +++ b/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): diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index b01a50e2fa..168278cb4f 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2179,6 +2179,51 @@ 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 + ) # Use FAST_COMPILE to prevent optimizations from changing the graph structure + + # Find the Join op node in the graph + join_nodes = [ + node for node in f.maker.fgraph.toposort() if isinstance(node.op, Join) + ] + assert len(join_nodes) == 1, "Expected exactly one Join node in the graph" + join_node = join_nodes[0] + + # 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_nodes = [ + node for node in f2.maker.fgraph.toposort() if isinstance(node.op, Join) + ] + assert len(join_nodes) == 1, "Expected exactly one Join node in the graph" + join_node = join_nodes[0] + + # 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) From a8de454f95f667877087c5d800f9fe9498e12fa3 Mon Sep 17 00:00:00 2001 From: Pablo de Roque Date: Tue, 8 Jul 2025 11:21:08 +0200 Subject: [PATCH 2/7] Update tests/tensor/test_basic.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/tensor/test_basic.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 168278cb4f..f48a232b1b 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2188,17 +2188,7 @@ def test_join_negative_axis_rewrite(self): # 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 - ) # Use FAST_COMPILE to prevent optimizations from changing the graph structure - - # Find the Join op node in the graph - join_nodes = [ - node for node in f.maker.fgraph.toposort() if isinstance(node.op, Join) - ] - assert len(join_nodes) == 1, "Expected exactly one Join node in the graph" - join_node = join_nodes[0] + assert s.owner.op.axis == 1 # Check that the axis input has been converted to a constant with value 1 (not -1) axis_input = join_node.inputs[0] From ccd812dd18ab9796383e7d0d2c666c22289200e9 Mon Sep 17 00:00:00 2001 From: PabloRoque Date: Tue, 8 Jul 2025 12:37:55 +0200 Subject: [PATCH 3/7] Simplify test --- tests/tensor/test_basic.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index f48a232b1b..4d6062bc6e 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2188,7 +2188,12 @@ def test_join_negative_axis_rewrite(self): # Create join with negative axis s = join(-1, a, b) - assert s.owner.op.axis == 1 + # 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] @@ -2201,11 +2206,8 @@ def test_join_negative_axis_rewrite(self): s2 = join(-2, a, b) f2 = pytensor.function([], [s2], mode=self.mode) - join_nodes = [ - node for node in f2.maker.fgraph.toposort() if isinstance(node.op, Join) - ] - assert len(join_nodes) == 1, "Expected exactly one Join node in the graph" - join_node = join_nodes[0] + 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] From 0b23a66ec16cc8e851415e0fb384b3a936fb424a Mon Sep 17 00:00:00 2001 From: PabloRoque Date: Tue, 8 Jul 2025 14:49:02 +0200 Subject: [PATCH 4/7] Test without compiling functions --- tests/tensor/test_basic.py | 34 ++++++++++++++++------------------ 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 4d6062bc6e..13e5096682 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2188,33 +2188,31 @@ def test_join_negative_axis_rewrite(self): # 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) + assert isinstance( + s.owner.outputs[0].owner.op, Join + ), "Expected output node to be a Join op" - # 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 isinstance( + s.owner.inputs[0], 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}" + s.owner.inputs[0].data == 1 + ), f"Expected axis to be normalized to 1, got {s.owner.inputs[0].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" + assert isinstance( + s2.owner.outputs[0].owner.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 isinstance( + s2.owner.inputs[0], 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}" + s2.owner.inputs[0].data == 0 + ), f"Expected axis to be normalized to 0, got {s2.owner.inputs[0].data}" def test_TensorFromScalar(): From 8657a642f521fd2e4f99d80fa20e1fba70a8aa2b Mon Sep 17 00:00:00 2001 From: Pablo de Roque Date: Tue, 8 Jul 2025 15:39:59 +0200 Subject: [PATCH 5/7] Update tests/tensor/test_basic.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/tensor/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 13e5096682..32841a1355 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2180,7 +2180,7 @@ def test_join_performance(self, ndim, axis, memory_layout, gc, benchmark): benchmark(fn, *test_values) def test_join_negative_axis_rewrite(self): - """Test that constant negative axis is rewritten to positive axis during canonicalization.""" + """Test that constant negative axis is rewritten to positive axis in make_node.""" 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) From 26071f7dfefdd67d6ea808046d5c38dfb49f5fcc Mon Sep 17 00:00:00 2001 From: PabloRoque Date: Tue, 8 Jul 2025 16:17:25 +0200 Subject: [PATCH 6/7] Test using equal_computations --- tests/tensor/test_basic.py | 30 ++---------------------------- 1 file changed, 2 insertions(+), 28 deletions(-) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index 32841a1355..b0278d43fd 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2185,34 +2185,8 @@ def test_join_negative_axis_rewrite(self): a = self.shared(v) b = as_tensor_variable(v) - # Create join with negative axis - s = join(-1, a, b) - - assert isinstance( - s.owner.outputs[0].owner.op, Join - ), "Expected output node to be a Join op" - - assert isinstance( - s.owner.inputs[0], ptb.Constant - ), "Expected axis to be a Constant" - assert ( - s.owner.inputs[0].data == 1 - ), f"Expected axis to be normalized to 1, got {s.owner.inputs[0].data}" - - # Now test with axis -2 which should be rewritten to 0 - s2 = join(-2, a, b) - - assert isinstance( - s2.owner.outputs[0].owner.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) - assert isinstance( - s2.owner.inputs[0], ptb.Constant - ), "Expected axis to be a Constant" - assert ( - s2.owner.inputs[0].data == 0 - ), f"Expected axis to be normalized to 0, got {s2.owner.inputs[0].data}" + equal_computations([join(-1, a, b)], [join(1, a, b)]) + equal_computations([join(-2, a, b)], [join(0, a, b)]) def test_TensorFromScalar(): From 1071abc24bfd2bc9079c232ecabdb3a6657aa366 Mon Sep 17 00:00:00 2001 From: Pablo de Roque Date: Tue, 8 Jul 2025 16:39:20 +0200 Subject: [PATCH 7/7] Update tests/tensor/test_basic.py Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com> --- tests/tensor/test_basic.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/tensor/test_basic.py b/tests/tensor/test_basic.py index b0278d43fd..f3b68f0e14 100644 --- a/tests/tensor/test_basic.py +++ b/tests/tensor/test_basic.py @@ -2185,8 +2185,8 @@ def test_join_negative_axis_rewrite(self): a = self.shared(v) b = as_tensor_variable(v) - equal_computations([join(-1, a, b)], [join(1, a, b)]) - equal_computations([join(-2, a, b)], [join(0, a, b)]) + assert equal_computations([join(-1, a, b)], [join(1, a, b)]) + assert equal_computations([join(-2, a, b)], [join(0, a, b)]) def test_TensorFromScalar():