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

Code before:


@pytest.mark.skip_xp_backends(np_only=True)
def test_sparse_matmat_exception():
    B = sparse.eye_array(2)
    # well defined matmat via `aslinearoperator`
    A = interface.aslinearoperator(sparse.eye_array(2))
    assert isinstance(A @ B, sparse.sparray)
    assert isinstance(B @ A, sparse.sparray)
    xp_assert_equal((A @ B).toarray(), np.eye(2))
    xp_assert_equal((B @ A).toarray(), np.eye(2))
    # ill-defined matmat via default fallback to matvec
    A = interface.LinearOperator((2, 2), matvec=lambda x: x)
    msg = "Try wrapping the matrix with `aslinearoperator` first."
    with assert_raises(TypeError, match=msg):
        A @ B
    with assert_raises(TypeError, match=msg):
        B @ A
    # after using `aslinearoperator`
    B = interface.aslinearoperator(B)
    assert isinstance(A @ B, interface.LinearOperator)
    assert isinstance(B @ A, interface.LinearOperator)
    xp_assert_equal((A @ B).matvec(np.ones(2)), np.ones(2))


@pytest.mark.skip_xp_backends(np_only=True) # remaining reference with JAX
def test_MatrixLinearOperator_refcycle(xp):
    # gh-10634
    # Test that MatrixLinearOperator can be automatically garbage collected
    A = xp.eye(2)
    with assert_deallocated(interface.MatrixLinearOperator, A, xp) as op:
        op.adjoint()
        del op


@pytest.mark.parametrize("left", [False, True])
@pytest.mark.parametrize("operator_definition", ["subclass_matvec", "subclass_matmat",
                                                 "__init__matvec", "__init__matmat",
                                                 "aslinearoperator"])
@pytest.mark.parametrize("batch_A", [(), (5,), (0,)])
@pytest.mark.parametrize("batch_x", [(), (6, 1), (0, 1)])
@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
@pytest.mark.skip_xp_backends("dask.array", reason="https://github.com/dask/dask/issues/11711")
@pytest.mark.skip_xp_backends(
    "array_api_strict",
    reason="https://github.com/data-apis/array-api-strict/issues/188"
)
def test_batch(left, operator_definition, batch_A, batch_x, dtype, xp):
    # TODO ideas:
    # - test lower-precision types
    # - test `transpose`, `adjoint`, `__mul__`, etc.

Review comment (suggestion):
```
ref        = <weakref at 0x17d307510; to 'scipy.sparse.linalg._interface.MatrixLinearOperator' at 0x17d3e79d0>
```

Code after:
    xp_assert_equal(A @ X2, Y @ X2)
    xp_assert_equal(X2.mT @ A, X2.mT @ Y)
    xp_assert_equal(A.T @ X2, Y.mT @ X2)


@pytest.mark.skip_xp_backends(np_only=True)
def test_sparse_matmat_exception():
    B = sparse.eye_array(2)
    # well defined matmat via `aslinearoperator`
    A = interface.aslinearoperator(sparse.eye_array(2))
    assert isinstance(A @ B, sparse.sparray)
    assert isinstance(B @ A, sparse.sparray)
    xp_assert_equal((A @ B).toarray(), np.eye(2))
    xp_assert_equal((B @ A).toarray(), np.eye(2))
    # ill-defined matmat via default fallback to matvec
    A = interface.LinearOperator((2, 2), matvec=lambda x: x)
    msg = "Try wrapping the matrix with `aslinearoperator` first."
    with assert_raises(TypeError, match=msg):
        A @ B
    with assert_raises(TypeError, match=msg):
        B @ A
    # after using `aslinearoperator`
    B = interface.aslinearoperator(B)
    assert isinstance(A @ B, interface.LinearOperator)
    assert isinstance(B @ A, interface.LinearOperator)
    xp_assert_equal((A @ B).matvec(np.ones(2)), np.ones(2))


@pytest.mark.skip_xp_backends(
    "jax.numpy", reason="remaining reference (due to lazy_xp_function?)"
)
def test_MatrixLinearOperator_refcycle(xp):
    # gh-10634
    # Test that MatrixLinearOperator can be automatically garbage collected
    A = xp.eye(2)
    with assert_deallocated(interface.MatrixLinearOperator, A, xp) as op:
        op.adjoint()
        del op


@pytest.mark.parametrize("left", [False, True])
@pytest.mark.parametrize("operator_definition", ["subclass_matvec", "subclass_matmat",
                                                 "__init__matvec", "__init__matmat",
                                                 "aslinearoperator"])
@pytest.mark.parametrize("batch_A", [(), (5,), (0,)])
@pytest.mark.parametrize("batch_x", [(), (6, 1), (0, 1)])
@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
@pytest.mark.skip_xp_backends("dask.array", reason="https://github.com/dask/dask/issues/11711")
@pytest.mark.skip_xp_backends(
    "array_api_strict",
    reason="https://github.com/data-apis/array-api-strict/issues/188"