Does cotengra support sparse tensors, like convolutions?

[thomasahle]

The Hayashi et al. way to write a convolutional layer in a NN is,
einsum('c h w, h h2 i, w w2 j, c i j c2 -> c2 h2 w2', X, C, C, T), where C is a sparse tensor defined by C_{i,j,k} = [i=j+k]. (See also the image below.)

Obviously, directly materializing C this way would be very inefficient. It would turn an O(n^2) operation into n^3.

It would be cool if cotengra could optimize it though, and contract it using an efficient algorithm. Besides allowing simple notation for CNNs, it would make it easier to do research like Felix Dangel's paper Convolutions and More as Einsum exploring new architectures.

[jcmgray]

Hi @thomasahle, in terms of explicit or symbolic representation of this kind of tensor and the reduced cost of operating on it, sadly no. The cost when optimizing the contraction tree is always assumed to be like a dense tensor. I would think that this is probably still a reasonable heuristic, given, at least for large networks, the most important thing is buildup of many dimensions.

In terms of performing the contraction cotengra supports anything with with (batch) matmul, transpose, and reshape (or preferably its own tensordot and einsum impl) e.g. https://sparse.pydata.org/en/stable/index.html. I believe pydata sparse arrays contract with dense numpy arrays fine, so it should just be a matter of constructing the sparse form of $C$ and supplying it as a standard input.

[thomasahle]

Hm, you might be right about memory.
For performing the contraction, are you saying I can just write a class like this?

class SparseConv1D:
    def __init__(self, I, J, K):
        self.I, self.J, self.K = I, J, K

    @property
    def shape(self):
        return (self.I, self.J, self.K)

    def __matmul__(self, x):
        """
        If x has shape (J, K, ...) then (C @ x) has shape (I, ...).
        The contraction sums over j,k where i == j+k.
        """
        if x.shape[0] != self.J or x.shape[1] != self.K:
            raise ValueError(f"x must have shape ({self.J},{self.K},...)")

        out = np.zeros((self.I,) + x.shape[2:], dtype=x.dtype)
        for i in range(self.I):
            # For each i, we sum x[j,k,...] for all j,k s.t. i == j+k.
            # That implies j = i - k, if valid in [0,J).
            for k in range(self.K):
                j = i - k
                if 0 <= j < self.J:
                    out[i] += x[j, k]
        return out

    def __repr__(self):
        return f"SparseConv1D(shape={self.shape})"

[jcmgray]

If you wanted to implement your own custom class, rather than just use sparse, you would have to either

• A): implement transpose, reshape, and batch matmul (i.e. signature "Bij,Bjk->Bik") but this would need to work on any grouping of the different axes. Non batch matmul would also be fine if the contractions have no hyper indices
• B): implement pairwise einsum directly (or alternatively just tensordot if the contractions have no hyper indices).

In the first case you would call something like:

# use cotengra's implementation of `tensordot` and `einsum`
contract(..., implementation="cotengra")

In the second case you would call something like:

# only necessary if your class is not defined in the `myconvlib` module
# or `einsum` is not defined in the top level namespace of that module
import autoray as ar
ar.register_backend(SparseConv1D, "myconvlib")
ar.register_function("myconvlib", "einsum", my_einsum_impl)
ar.register_function("myconvlib", "tensordot", my_tensordot_impl)

# then contract as usual, or if you did not provide tensordot:
contract(..., prefer_einsum=True)

I would guess using an existing sparse implementation is easiest however.

[thomasahle]

When I use sparse.COO I get

        if out is None:
>           return _VF.tensordot(a, b, dims_a, dims_b)  # type: ignore[attr-defined]
E           TypeError: tensordot(): argument 'other' (position 2) must be Tensor, not COO

/opt/homebrew/lib/python3.12/site-packages/torch/functional.py:1355: TypeError

When I use sparse.CSR I get

>           return _VF.tensordot(a, b, dims_a, dims_b)  # type: ignore[attr-defined]
E           RuntimeError: reshape is not implemented for sparse tensors

/opt/homebrew/lib/python3.12/site-packages/torch/functional.py:1355: RuntimeError

When I use torch.sparse_csr_tensor I get

E           NotImplementedError: Could not run 'aten::as_strided' with arguments from the 'SparseCsrCPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build).

The code looks like this: https://gist.github.com/thomasahle/f5e8c38c8adfaa5a7da58ba46c21d7b0

I also tried all three sparse formats with torch.einsum, and it doens't work either.

[jcmgray]

Ah yes I'm afraid I'm not really aware of the state of things for torch. sparse is compatible with numpy, in a pinch you could define the pairwise contraction to convert from torch -> numpy, do the contraction and convert back, would probably be okay on the cpu but not very satisfying.

[thomasahle]

I ended up writing this terrible mess, enumerating all the different patterns in which a convolution might connect with another tensor, and applying fold/unfold or other functions specially written for that case.

Mostly because I thought it would be too expensive to switch from torch to numpy... But maybe it was a mistake.

How fast would you generally assume dense-sparse multiplication to be compared to fold/unfold?

[jcmgray]

Mostly because I thought it would be too expensive to switch from torch to numpy... But maybe it was a mistake.

I think numpy <-> torch cpu can be zero copy and thus pretty instant, but for CUDA and autograd there may be complications.

How fast would you generally assume dense-sparse multiplication to be compared to fold/unfold?

I'm not totally sure but some loose thoughts:

I think a very common way to do sparse contractions is anyway to fuse/fold the op into matrix multiplication so one can use e.g. a matmul-suited sparse format like CSR.

I don't use sparse-dense multiplication much, but I'd assume for each sparse dimension with size $d$ average sparsity $k$ the speedup to scale like $d/k$ but with some pretty large constant overhead.

I'm not super familiar with the fold/unfold term, my impression is that it is a type of transposing and reshaping ('fusing') right? So essentially linear in the number of entries. So it will always scale at least as well as the contraction but practically such operations can still be significant, especially for memory bound contractions.