Repository: scipy/scipy
File: scipy/sparse/linalg/tests/test_interface.py (Python)

Code before:
                               A_ @ x_col if left else A_.mT.conj() @ x_col)
    # c. with matrix (or batch of matrices)
    np.testing.assert_allclose(A.matmat(x_mat) if left else A.rmatmat(x_mat),
                               A_ @ x_mat if left else A_.mT.conj() @ x_mat)

    # test __matmul__ (via `@`), `__call__`, and `dot`
    if left:
        # a. with row vector (or batch of row vectors)
        if batch_x:
            with pytest.raises(ValueError, match="Dimension mismatch:..."):
                A @ x_row
        else:
            np.testing.assert_allclose(A @ x_row, A_ @ x_row)
            np.testing.assert_allclose(A(x_row), A_ @ x_row)
            np.testing.assert_allclose(A.dot(x_row), A_ @ x_row)
        # b. with column vector (or batch of column vectors)
        np.testing.assert_allclose(A @ x_col, A_ @ x_col)
        np.testing.assert_allclose(A(x_col), A_ @ x_col)
        np.testing.assert_allclose(A.dot(x_col), A_ @ x_col)
        # c. with matrix (or batch of matrices)
        np.testing.assert_allclose(A @ x_mat, A_ @ x_mat)
        np.testing.assert_allclose(A(x_mat), A_ @ x_mat)
        np.testing.assert_allclose(A.dot(x_mat), A_ @ x_mat)
    else:
        # a. with row vector (or batch of row vectors)
        if batch_x and batch_x[-1] != 1: # interpreted as column vector
            if batch_A:
                # Maybe should be "Dimension mismatch..."?
                msg = "operands could not be broadcast..."
                with pytest.raises(ValueError, match=msg):
                    x_row @ A
            else:
                np.testing.assert_allclose(x_row @ A, x_row @ A_)
                np.testing.assert_allclose(A.rdot(x_row), x_row @ A_)
        # b. with column vector (or batch of column vectors)
        with pytest.raises(ValueError, match="Dimension mismatch:..."):
            x_col @ A
        # c. with matrix (or batch of matrices)
        np.testing.assert_allclose(x_mat.mT.conj() @ A, x_mat.mT.conj() @ A_)
        np.testing.assert_allclose(A.rdot(x_mat.mT.conj()), x_mat.mT.conj() @ A_)


Review comment (performance):
Sorry, could you elaborate? Doesn't this effectively skip this part of the tests, since all cases in which `batch_x` is truthy have `batch_x[-1] == 1`?

Code after:
                               A_ @ x_col if left else A_.mT.conj() @ x_col)
    # c. with matrix (or batch of matrices)
    np.testing.assert_allclose(A.matmat(x_mat) if left else A.rmatmat(x_mat),
                               A_ @ x_mat if left else A_.mT.conj() @ x_mat)

    # test __matmul__ (via `@`), `__call__`, and `dot`
    if left:
        # a. with row vector (or batch of row vectors)
        if batch_x:
            with pytest.raises(ValueError, match="Dimension mismatch:..."):
                A @ x_row
        else:
            np.testing.assert_allclose(A @ x_row, A_ @ x_row)
            np.testing.assert_allclose(A(x_row), A_ @ x_row)
            np.testing.assert_allclose(A.dot(x_row), A_ @ x_row)
        # b. with column vector (or batch of column vectors)
        np.testing.assert_allclose(A @ x_col, A_ @ x_col)
        np.testing.assert_allclose(A(x_col), A_ @ x_col)
        np.testing.assert_allclose(A.dot(x_col), A_ @ x_col)
        # c. with matrix (or batch of matrices)
        np.testing.assert_allclose(A @ x_mat, A_ @ x_mat)
        np.testing.assert_allclose(A(x_mat), A_ @ x_mat)
        np.testing.assert_allclose(A.dot(x_mat), A_ @ x_mat)
    else:
        # a. with row vector (or batch of row vectors)
        broadcastable = True
        try:
            # `batch_x[-1]` becomes part of the core shape for matrix
            # multiplication with 1-D row vector
            # success thus depends on whether `batch_A` can broadcast
            # with `batch_x[:-1]`
            np.broadcast_shapes(batch_A, batch_x[:-1])
        except ValueError:
            broadcastable = False
        if not broadcastable:
            # Maybe should be "Dimension mismatch..."?
            msg = "operands could not be broadcast..."
            with pytest.raises(ValueError, match=msg):
                x_row @ A
        else:
            np.testing.assert_allclose(x_row @ A, x_row @ A_)
            np.testing.assert_allclose(A.rdot(x_row), x_row @ A_)
        # b. with column vector (or batch of column vectors)
        with pytest.raises(ValueError, match="Dimension mismatch:..."):
            x_col @ A
        # c. with matrix (or batch of matrices)
        np.testing.assert_allclose(x_mat.mT.conj() @ A, x_mat.mT.conj() @ A_)
        np.testing.assert_allclose(A.rdot(x_mat.mT.conj()), x_mat.mT.conj() @ A_)