pair coding results

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -162,59 +162,40 @@ def local_block_diag_dot_to_dot_block_diag(fgraph, node):
     if not isinstance(node.op.core_op, BlockDiagonal):
         return
 
-    def check_for_block_diag(x):
-        return x.owner and (
-            isinstance(x.owner.op, BlockDiagonal)
-            or isinstance(x.owner.op, Blockwise)
-            and isinstance(x.owner.op.core_op, BlockDiagonal)
-        )
-
     # Check that the BlockDiagonal is an input to a Dot node:
     for client in get_clients_at_depth(fgraph, node, depth=1):
-        if not isinstance(client.op, Dot):
+        if not (
+            (
+                isinstance(client.op, Dot)
+                and all(input.ndim == 2 for input in client.inputs)
+            )
+            or client.op == _matrix_matrix_matmul
+        ):
             continue
 
         op = client.op
-        x, y = client.inputs
 
-        if not (check_for_block_diag(x) or check_for_block_diag(y)):
-            continue
+        client_idx = client.inputs.index(node.outputs[0])
 
-        # Case 1: Only one input is BlockDiagonal. In this case, multiply all components of the block-diagonal with the
-        # non-block diagonal, and return a new block diagonal
-        if check_for_block_diag(x) and not check_for_block_diag(y):
-            components = x.owner.inputs
-            y_splits = split(
-                y,
-                splits_size=[component.shape[-1] for component in components],
-                n_splits=len(components),
-            )
-            new_components = [
-                op(component, y_split)
-                for component, y_split in zip(components, y_splits)
-            ]
-            new_output = join(0, *new_components)
-
-        elif not check_for_block_diag(x) and check_for_block_diag(y):
-            components = y.owner.inputs
-            x_splits = split(
-                x,
-                splits_size=[component.shape[0] for component in components],
-                n_splits=len(components),
-                axis=1,
-            )
+        other_input = client.inputs[1 - client_idx]
+        components = node.inputs
 
-            new_components = [
-                op(x_split, component)
-                for component, x_split in zip(components, x_splits)
-            ]
-            new_output = join(1, *new_components)
+        split_axis = -2 if client_idx == 0 else -1
+        shape_idx = -1 if client_idx == 0 else -2
 
-        # Case 2: Both inputs are BlockDiagonal. Do nothing
-        else:
-            # TODO: If shapes are statically known and all components have equal shapes, we could rewrite
-            #  this case to block_diag(*[dot(comp_1, comp_2) for comp_1, comp_2 in zip(x.owner.inputs, y.owner.inputs)])
-            continue
+        other_dot_input_split = split(
+            other_input,
+            splits_size=[component.shape[shape_idx] for component in components],
+            n_splits=len(components),
+            axis=split_axis,
+        )
+        new_components = [
+            op(component, other_split)
+            if client_idx == 0
+            else op(other_split, component)
+            for component, other_split in zip(components, other_dot_input_split)
+        ]
+        new_output = join(split_axis, *new_components)
 
         copy_stack_trace(node.outputs[0], new_output)
         return {client.outputs[0]: new_output}

tests/tensor/rewriting/test_math.py

@@ -4749,15 +4749,21 @@ def test_local_dot_to_mul(batched, a_shape, b_shape):
 
 
 @pytest.mark.parametrize("left_multiply", [True, False], ids=["left", "right"])
-def test_local_block_diag_dot_to_dot_block_diag(left_multiply):
+@pytest.mark.parametrize(
+    "batch_left", [True, False], ids=["batched_left", "unbatched_left"]
+)
+@pytest.mark.parametrize(
+    "batch_right", [True, False], ids=["batched_right", "unbatched_right"]
+)
+def test_local_block_diag_dot_to_dot_block_diag(left_multiply, batch_left, batch_right):
     """
     Test that dot(block_diag(x, y,), z) is rewritten to concat(dot(x, z[:n]), dot(y, z[n:]))
     """
     a = tensor("a", shape=(4, 2))
-    b = tensor("b", shape=(2, 4))
+    b = tensor("b", shape=(2, 4) if not batch_left else (3, 2, 4))
     c = tensor("c", shape=(4, 4))
     d = tensor("d", shape=(10, 10))
-    e = tensor("e", shape=(10, 10))
+    e = tensor("e", shape=(10, 10) if not batch_right else (3, 1, 10, 10))
 
     x = pt.linalg.block_diag(a, b, c)
 
@@ -4767,30 +4773,38 @@ def test_local_block_diag_dot_to_dot_block_diag(left_multiply):
     else:
         out = [d @ x, e @ x]
 
-    fn = pytensor.function([a, b, c, d, e], out, mode=rewrite_mode)
+    with config.change_flags(optimizer_verbose=True):
+        fn = pytensor.function([a, b, c, d, e], out, mode=rewrite_mode)
+
     assert not any(
         isinstance(node.op, BlockDiagonal) for node in fn.maker.fgraph.toposort()
     )
 
     fn_expected = pytensor.function(
         [a, b, c, d, e],
         out,
-        mode=rewrite_mode.excluding("local_block_diag_dot_to_dot_block_diag"),
+        mode=Mode(linker="py", optimizer=None),
     )
 
+    # TODO: Count Dots
+
     rng = np.random.default_rng()
     a_val = rng.normal(size=a.type.shape).astype(a.type.dtype)
     b_val = rng.normal(size=b.type.shape).astype(b.type.dtype)
     c_val = rng.normal(size=c.type.shape).astype(c.type.dtype)
     d_val = rng.normal(size=d.type.shape).astype(d.type.dtype)
     e_val = rng.normal(size=e.type.shape).astype(e.type.dtype)
 
-    np.testing.assert_allclose(
-        fn(a_val, b_val, c_val, d_val, e_val),
-        fn_expected(a_val, b_val, c_val, d_val, e_val),
-        atol=1e-6 if config.floatX == "float32" else 1e-12,
-        rtol=1e-6 if config.floatX == "float32" else 1e-12,
-    )
+    rewrite_outs = fn(a_val, b_val, c_val, d_val, e_val)
+    expected_outs = fn_expected(a_val, b_val, c_val, d_val, e_val)
+
+    for out, expected in zip(rewrite_outs, expected_outs):
+        np.testing.assert_allclose(
+            out,
+            expected,
+            atol=1e-6 if config.floatX == "float32" else 1e-12,
+            rtol=1e-6 if config.floatX == "float32" else 1e-12,
+        )
 
 
 @pytest.mark.parametrize("rewrite", [True, False], ids=["rewrite", "no_rewrite"])