Sparse and opt-einsum fail

sparse_einsum_fail.py

import numpy as np
from opt_einsum import contract, helpers, contract_expression
import sparse
import pytest


tests = [
    'ab,bc->ca',
    'abc,bcd,dea',
    'abc,def->fedcba',
    # test 'prefer einsum' ops
    'ijk,ikj',
    'i,j->ij',
    'ijk,k->ij',
    'ac,bc->abc',
    'aci,bci->abci',
]


@pytest.mark.parametrize("string", tests)
def test_sparse(string):
    views = helpers.build_views(string)

    # sparsify views so they don't become dense during contraction
    for view in views:
        np.random.seed(42)
        mask = np.random.choice([False, True], view.shape, True, [0.05, 0.95])
        view[mask] = 0

    # test with NumPy arrays
    ein = contract(string, *views, optimize=False, use_blas=False)
    numpy_ein = np.einsum(string, *views)
    assert np.allclose(numpy_ein, ein)

    sparse_views = [sparse.COO.from_numpy(x) for x in views]
    # pass SparseArrays to contract

    if string in {'ac,bc->abc', 'aci,bci->abci'}:
        with pytest.raises(AttributeError, match=("sparse doesn't seem to provide "
                                                  "the function einsum")):
            contract(string, *sparse_views, backend='sparse')
    else:
        sparse_opt = contract(string, *sparse_views, backend='sparse')
        assert np.allclose(ein, sparse_opt.todense())