Address PR feedback for matrix-vector operations

- Improve docstrings with concrete shape examples
- Explicitly state equivalence to NumPy functions
- Simplify tests into a single parametrized test
- Add dtype parameter test to ensure full coverage
- Keep implementation minimal by relying on Blockwise checks

🤖 Generated with Claude Code
Co-Authored-By: Claude <[email protected]>

Author: twiecki

pytensor/tensor/math.py

@@ -2864,17 +2864,24 @@ def vecdot(
 
     Notes
     -----
-    This is similar to `np.vecdot` and computes the dot product of
+    This is equivalent to `numpy.vecdot` and computes the dot product of
     vectors along the last axis of both inputs. Broadcasting is supported
     across all other dimensions.
 
     Examples
     --------
     >>> import pytensor.tensor as pt
-    >>> x = pt.matrix("x")
-    >>> y = pt.matrix("y")
-    >>> z = pt.vecdot(x, y)
-    >>> # Equivalent to np.sum(x * y, axis=-1)
+    >>> # Vector dot product with shape (5,) inputs
+    >>> x = pt.vector("x")  # shape (5,)
+    >>> y = pt.vector("y")  # shape (5,)
+    >>> z = pt.vecdot(x, y)  # scalar output
+    >>> # Equivalent to numpy.vecdot(x, y) or numpy.sum(x * y)
+    >>>
+    >>> # With batched inputs of shape (3, 5)
+    >>> x_batch = pt.matrix("x")  # shape (3, 5)
+    >>> y_batch = pt.matrix("y")  # shape (3, 5)
+    >>> z_batch = pt.vecdot(x_batch, y_batch)  # shape (3,)
+    >>> # Equivalent to numpy.sum(x_batch * y_batch, axis=-1)
     """
     x1 = as_tensor_variable(x1)
     x2 = as_tensor_variable(x2)
@@ -2918,17 +2925,16 @@ def matvec(
     Examples
     --------
     >>> import pytensor.tensor as pt
-    >>> import numpy as np
     >>> # Matrix-vector product
-    >>> A = pt.matrix("A")  # shape (M, K)
-    >>> v = pt.vector("v")  # shape (K,)
-    >>> result = pt.matvec(A, v)  # shape (M,)
-    >>> # Equivalent to np.matmul(A, v)
+    >>> A = pt.matrix("A")  # shape (3, 4)
+    >>> v = pt.vector("v")  # shape (4,)
+    >>> result = pt.matvec(A, v)  # shape (3,)
+    >>> # Equivalent to numpy.matmul(A, v)
     >>>
     >>> # Batched matrix-vector product
-    >>> batched_A = pt.tensor3("A")  # shape (B, M, K)
-    >>> batched_v = pt.matrix("v")  # shape (B, K)
-    >>> result = pt.matvec(batched_A, batched_v)  # shape (B, M)
+    >>> batched_A = pt.tensor3("A")  # shape (2, 3, 4)
+    >>> batched_v = pt.matrix("v")  # shape (2, 4)
+    >>> result = pt.matvec(batched_A, batched_v)  # shape (2, 3)
     """
     x1 = as_tensor_variable(x1)
     x2 = as_tensor_variable(x2)
@@ -2971,17 +2977,16 @@ def vecmat(
     Examples
     --------
     >>> import pytensor.tensor as pt
-    >>> import numpy as np
     >>> # Vector-matrix product
-    >>> v = pt.vector("v")  # shape (K,)
-    >>> A = pt.matrix("A")  # shape (K, N)
-    >>> result = pt.vecmat(v, A)  # shape (N,)
-    >>> # Equivalent to np.matmul(v, A)
+    >>> v = pt.vector("v")  # shape (3,)
+    >>> A = pt.matrix("A")  # shape (3, 4)
+    >>> result = pt.vecmat(v, A)  # shape (4,)
+    >>> # Equivalent to numpy.matmul(v, A)
     >>>
     >>> # Batched vector-matrix product
-    >>> batched_v = pt.matrix("v")  # shape (B, K)
-    >>> batched_A = pt.tensor3("A")  # shape (B, K, N)
-    >>> result = pt.vecmat(batched_v, batched_A)  # shape (B, N)
+    >>> batched_v = pt.matrix("v")  # shape (2, 3)
+    >>> batched_A = pt.tensor3("A")  # shape (2, 3, 4)
+    >>> result = pt.vecmat(batched_v, batched_A)  # shape (2, 4)
     """
     x1 = as_tensor_variable(x1)
     x2 = as_tensor_variable(x2)

tests/tensor/test_math.py

@@ -2080,55 +2080,36 @@ def is_super_shape(var1, var2):
 
 
 class TestMatrixVectorOps:
-    def test_vecdot(self):
-        """Test vecdot function with various input shapes."""
-        rng = np.random.default_rng(seed=utt.fetch_seed())
-
-        # Test vector-vector
-        x = vector()
-        y = vector()
-        z = vecdot(x, y)
-        f = function([x, y], z)
-        x_val = random(5, rng=rng).astype(config.floatX)
-        y_val = random(5, rng=rng).astype(config.floatX)
-        np.testing.assert_allclose(f(x_val, y_val), np.dot(x_val, y_val))
-
-        # Test batched vectors
-        x = tensor3()
-        y = tensor3()
-        z = vecdot(x, y)
-        f = function([x, y], z)
-
-        x_val = random(2, 3, 4, rng=rng).astype(config.floatX)
-        y_val = random(2, 3, 4, rng=rng).astype(config.floatX)
-        expected = np.sum(x_val * y_val, axis=-1)
-        np.testing.assert_allclose(f(x_val, y_val), expected)
+    """Test vecdot, matvec, and vecmat helper functions."""
 
     @pytest.mark.parametrize(
-        "func,x_shape,y_shape,make_expected",
+        "func,x_shape,y_shape,np_func,batch_axis",
         [
-            # matvec tests - Matrix(M,K) @ Vector(K) -> Vector(M)
-            (matvec, (3, 4), (4,), lambda x, y: np.dot(x, y)),
-            # matvec batch tests - Tensor3(B,M,K) @ Matrix(B,K) -> Matrix(B,M)
+            # vecdot
+            (vecdot, (5,), (5,), lambda x, y: np.dot(x, y), None),
+            (vecdot, (3, 5), (3, 5), lambda x, y: np.sum(x * y, axis=-1), -1),
+            # matvec
+            (matvec, (3, 4), (4,), lambda x, y: np.dot(x, y), None),
             (
                 matvec,
                 (2, 3, 4),
                 (2, 4),
                 lambda x, y: np.array([np.dot(x[i], y[i]) for i in range(len(x))]),
+                0,
             ),
-            # vecmat tests - Vector(K) @ Matrix(K,N) -> Vector(N)
-            (vecmat, (3,), (3, 4), lambda x, y: np.dot(x, y)),
-            # vecmat batch tests - Matrix(B,K) @ Tensor3(B,K,N) -> Matrix(B,N)
+            # vecmat
+            (vecmat, (3,), (3, 4), lambda x, y: np.dot(x, y), None),
             (
                 vecmat,
                 (2, 3),
                 (2, 3, 4),
                 lambda x, y: np.array([np.dot(x[i], y[i]) for i in range(len(x))]),
+                0,
             ),
         ],
     )
-    def test_mat_vec_ops(self, func, x_shape, y_shape, make_expected):
-        """Parametrized test for matvec and vecmat functions."""
+    def test_matrix_vector_ops(self, func, x_shape, y_shape, np_func, batch_axis):
+        """Test all matrix-vector helper functions."""
         rng = np.random.default_rng(seed=utt.fetch_seed())
 
         # Create PyTensor variables with appropriate dimensions
@@ -2146,18 +2127,25 @@ def test_mat_vec_ops(self, func, x_shape, y_shape, make_expected):
         else:
             y = tensor3()
 
-        # Apply the function
+        # Test basic functionality
         z = func(x, y)
         f = function([x, y], z)
 
-        # Create random values
         x_val = random(*x_shape, rng=rng).astype(config.floatX)
         y_val = random(*y_shape, rng=rng).astype(config.floatX)
 
-        # Compare with the expected result
-        expected = make_expected(x_val, y_val)
+        expected = np_func(x_val, y_val)
         np.testing.assert_allclose(f(x_val, y_val), expected)
 
+        # Test with dtype parameter (to improve code coverage)
+        # Use float64 to ensure we can detect the difference
+        z_dtype = func(x, y, dtype="float64")
+        f_dtype = function([x, y], z_dtype)
+
+        result = f_dtype(x_val, y_val)
+        assert result.dtype == np.float64
+        np.testing.assert_allclose(result, expected)
+
 
 class TestTensordot:
     def TensorDot(self, axes):