Propagate triangularity through matrix transpose

pytensor/assumptions/triangular.py

@@ -136,21 +136,34 @@ def _lu_upper(key, op, feature, fgraph, node, input_states):
     return states
 
 
-def _dimshuffle_left_expand_dims(key, op, feature, fgraph, node, input_states):
-    """Triangularity survives left-expand-dims (batch broadcast).
-
-    Matrix transpose swaps lower<->upper, so it is *not* propagated under
-    the same key.
+@register_assumption(LOWER_TRIANGULAR, DimShuffle)
+def _lower_dimshuffle(key, op, feature, fgraph, node, input_states):
+    """Lower-triangularity survives left-expand-dims and arises from the matrix
+    transpose of an upper-triangular operand (since transposing swaps the two
+    triangles).
     """
-    if input_states[0] is not FactState.TRUE:
-        return [FactState.UNKNOWN]
-    if left_expand_dims_propagates_matrix_property(op):
+    if op.is_matrix_transpose:
+        return true_if(feature.check(node.inputs[0], UPPER_TRIANGULAR))
+    if input_states[
+        0
+    ] is FactState.TRUE and left_expand_dims_propagates_matrix_property(op):
         return [FactState.TRUE]
     return [FactState.UNKNOWN]
 
 
-register_assumption(LOWER_TRIANGULAR, DimShuffle)(_dimshuffle_left_expand_dims)
-register_assumption(UPPER_TRIANGULAR, DimShuffle)(_dimshuffle_left_expand_dims)
+@register_assumption(UPPER_TRIANGULAR, DimShuffle)
+def _upper_dimshuffle(key, op, feature, fgraph, node, input_states):
+    """Upper-triangularity survives left-expand-dims and arises from the matrix
+    transpose of a lower-triangular operand (since transposing swaps the two
+    triangles).
+    """
+    if op.is_matrix_transpose:
+        return true_if(feature.check(node.inputs[0], LOWER_TRIANGULAR))
+    if input_states[
+        0
+    ] is FactState.TRUE and left_expand_dims_propagates_matrix_property(op):
+        return [FactState.TRUE]
+    return [FactState.UNKNOWN]
 
 
 @register_assumption(LOWER_TRIANGULAR, Elemwise)

tests/assumptions/test_dimshuffle.py

@@ -38,15 +38,6 @@ def test_multiple_left_axes(self, key):
         _, af = make_fgraph(y)
         assert af.check(y, key)
 
-    def test_triangular_transpose_does_not_propagate(self):
-        """Transpose swaps the triangle, so a lower-triangular matrix is no longer
-        known to be lower-triangular."""
-        x = pt.matrix("x", shape=(4, 4))
-        x_lower = assume(x, lower_triangular=True)
-        y = x_lower.T
-        _, af = make_fgraph(y)
-        assert af.get(y, LOWER_TRIANGULAR) == FactState.UNKNOWN
-
     def test_right_expand_dims_does_not_propagate(self):
         """Adding a broadcast dim on the right shifts the matrix axes; the property is lost."""
         x = pt.matrix("x", shape=(4, 4))

tests/assumptions/test_triangular.py

@@ -166,3 +166,48 @@ def test_lu_permute_l_pl_is_not_lower_triangular():
     _, af = make_fgraph(PL, U)
     assert af.get(PL, LOWER_TRIANGULAR) == FactState.UNKNOWN
     assert af.check(U, UPPER_TRIANGULAR)
+
+
+class TestMatrixTransposeFlipsTriangle:
+    """Matrix transpose maps lower-triangular to upper-triangular and vice versa."""
+
+    @pytest.mark.parametrize(
+        "asserted, flipped",
+        [
+            (LOWER_TRIANGULAR, UPPER_TRIANGULAR),
+            (UPPER_TRIANGULAR, LOWER_TRIANGULAR),
+        ],
+    )
+    def test_transpose_flips_triangle(self, asserted, flipped):
+        x = pt.matrix("x", shape=(4, 4))
+        x_tagged = assume(x, **{asserted.name: True})
+        y = x_tagged.T
+        _, af = make_fgraph(y)
+        assert af.check(y, flipped)
+        assert af.get(y, asserted) == FactState.UNKNOWN
+
+    @pytest.mark.parametrize(
+        "asserted, flipped",
+        [
+            (LOWER_TRIANGULAR, UPPER_TRIANGULAR),
+            (UPPER_TRIANGULAR, LOWER_TRIANGULAR),
+        ],
+    )
+    def test_batched_transpose_flips_triangle(self, asserted, flipped):
+        x = pt.tensor("x", shape=(2, 3, 4, 4))
+        x_tagged = assume(x, **{asserted.name: True})
+        y = x_tagged.mT
+        _, af = make_fgraph(y)
+        assert af.check(y, flipped)
+        assert af.get(y, asserted) == FactState.UNKNOWN
+
+    @pytest.mark.parametrize("key", [LOWER_TRIANGULAR, UPPER_TRIANGULAR])
+    def test_double_transpose_recovers_triangle(self, key):
+        # Exercises the cross-key inference chain: the first transpose flips
+        # via a feature.check on the input, and the second transpose flips
+        # back via a feature.check on the freshly-inferred intermediate.
+        x = pt.matrix("x", shape=(4, 4))
+        x_tagged = assume(x, **{key.name: True})
+        y = x_tagged.T.T
+        _, af = make_fgraph(y)
+        assert af.check(y, key)