Handle symbolic shapes

Author: AllenDowney

pytensor/xtensor/rewriting/math.py

@@ -41,11 +41,11 @@ def lower_dot(fgraph, node):
     # Perform the einsum operation
     out_tensor = einsum(einsum_str, x_tensor, y_tensor)
 
-    # Reshape to match the expected output shape
-    try:
+    # Check if we have symbolic shapes
+    sym_shape = any(not isinstance(s, int) for s in out.type.shape)
+
+    # If we have concrete shapes, reshape to match them
+    if not sym_shape:
         out_tensor = reshape(out_tensor, out.type.shape)
-    except (TypeError, ValueError):
-        # Skip reshaping if symbolic shapes are present
-        pass
 
     return [xtensor_from_tensor(out_tensor, out.type.dims)]

tests/xtensor/test_math.py

@@ -277,6 +277,17 @@ def test_dot():
     z_test = fn(x_test, y_test)
     xr_assert_allclose(z_test, expected)
 
+    # Test symbolic shapes
+    x = xtensor("x", dims=("a", "b"), shape=(None, 3))  # First dimension is symbolic
+    y = xtensor("y", dims=("b", "c"), shape=(3, None))  # Second dimension is symbolic
+    z = x.dot(y)
+    fn = xr_function([x, y], z)
+    x_test = DataArray(np.ones((2, 3)), dims=("a", "b"))
+    y_test = DataArray(np.ones((3, 4)), dims=("b", "c"))
+    z_test = fn(x_test, y_test)
+    expected = x_test.dot(y_test)
+    xr_assert_allclose(z_test, expected)
+
 
 def test_dot_errors():
     x = xtensor("x", dims=("a", "b"), shape=(2, 3))