Preserve static shape information in block_diag

Author: jessegrabowski

pytensor/tensor/slinalg.py

@@ -1651,7 +1651,18 @@ class BlockDiagonal(BaseBlockDiagonal):
     def make_node(self, *matrices):
         matrices = self._validate_and_prepare_inputs(matrices, pt.as_tensor)
         dtype = _largest_common_dtype(matrices)
-        out_type = pytensor.tensor.matrix(dtype=dtype)
+
+        input_shapes = tuple(zip(*(m.type.shape for m in matrices)))
+        out_shape = tuple(
+            [
+                sum(dim_shapes)
+                if not any(shape is None for shape in dim_shapes)
+                else None
+                for dim_shapes in input_shapes
+            ]
+        )
+
+        out_type = pytensor.tensor.matrix(shape=out_shape, dtype=dtype)
         return Apply(self, matrices, [out_type])
 
     def perform(self, node, inputs, output_storage, params=None):

tests/tensor/test_slinalg.py

@@ -1040,10 +1040,20 @@ def test_block_diagonal():
     A = np.array([[1.0, 2.0], [3.0, 4.0]])
     B = np.array([[5.0, 6.0], [7.0, 8.0]])
     result = block_diag(A, B)
+    assert result.type.shape == (4, 4)
     assert result.owner.op.core_op._props_dict() == {"n_inputs": 2}
 
     np.testing.assert_allclose(result.eval(), scipy.linalg.block_diag(A, B))
 
+    A = pt.dmatrix("A", shape=(5, 5))
+    B = pt.dmatrix("B", shape=(3, None))
+    result = block_diag(A, B)
+    assert result.type.shape == (8, None)
+
+    A = pt.dmatrix("A", shape=(None, 5))
+    result = block_diag(A, B)
+    assert result.type.shape == (None, None)
+
 
 def test_block_diagonal_grad():
     A = np.array([[1.0, 2.0], [3.0, 4.0]])