How can one do a fast sparse tensor addition?

Below, I have the following code: We first generate 3 random sparse 1000x1000x1000 tensors with 10^6 entries each.

Then, I want to add them. But the usual way is incredibly slow. I mean with 10^6 entries, I expect runtimes of a few ms, but the usual way takes around 1s.

Some speedup by a factor 3 was possible by converting the 3-tensors to 2-tensors (i.e., 10^6x1000 matrices), then add, and then convert back (but via Partition, not via ArrayReshape). Does anybody know a faster way to do this? I mean, addition is such a basic operation that one would expect it to be fully optimized.

Cheers

RandomTensor[VV_,EE_]:=SparseArray[RandomInteger[{1, VV}, {EE, 3}] -> RandomReal[1,EE],{VV,VV,VV}];
VV=1000;
EE=1000000;
T1=RandomTensor[VV,EE];
T2=RandomTensor[VV,EE];
T3=RandomTensor[VV,EE];
Sum1=T1+T2+T3;//AbsoluteTiming
Sum2=Partition[ArrayFlatten[T1,1]+ArrayFlatten[T2,1]+ArrayFlatten[T3,1],VV]; //AbsoluteTiming
Max[Abs[Sum1-Sum2]]

• Hi, maybe you might want to call the numpy, pytorch, or tensorflow packages from Python within mathematica. The drawback is that requires knowing a bit of the Python language and that might take some time to set up. The following might depend on the version of mathematica but I personally would write the code on a jupyter notebook then copy paste the code into mathematica by typing > at the the beginning of a cell and then choosing Python. Oct 22, 2022 at 10:07

Mathematica uses a modified CSR format for SparseArrays of dimension two or higher. But the CSR format was actually only developed for matrices (dimension equal to 2). So it takes no wonder that it is also more efficient for matrices than for general tensors.

Hence converting to matrices is a good idea. You flattened the first two dimensions together, but I think flattening the last two dimensions is better for the following two reasons:

1. the column indices can be computed by a linear transformation, namely by T1["ColumnIndices"] . {{VV}, {1}}.

2. CSR does not like super long columns.

But maybe it is simply an implementation issue within Mathematica... CSR construction involves quite a lot of sorting, in particular, when one converts from the COO format. Here one can take several shortcuts for constructing the sparsity pattern of T1+T2, e.g., exploiting that the row pointers of T1 and T2 are already ordered and using a merge sort instead of first converting for COO and then back to CSR. However, these things are only worth doing in special applications. I guess that's why Wolfram Research took the easy track with COO as intermediate format (since the COO->CSR converter had been there already), at least for tensors of degree >2.

Anyways, apparently the conversion to and from matrices consumes a lot of time. So better always keep the data in this form. Here a timing comparison from my machine:

RandomTensor[VV_, EE_] := SparseArray[RandomInteger[{1, VV}, {EE, 3}] -> RandomReal[1, EE], {VV, VV, VV}];
First@AbsoluteTiming[
T1 = RandomTensor[VV, EE];
T2 = RandomTensor[VV, EE];
T3 = RandomTensor[VV, EE];
]
First@AbsoluteTiming[
Sum1 = T1 + T2 + T3;
]


0.271651

0.255094

vs. matrix-only operations:

RandomTensor2[VV_, EE_] :=  SparseArray[Transpose[{RandomInteger[{1, VV}, {EE}], RandomInteger[{1, VV VV}, {EE}]}] -> RandomReal[1, EE], {VV, VV VV}];
First@AbsoluteTiming[
S1 = RandomTensor2[VV, EE];
S2 = RandomTensor2[VV, EE];
S3 = RandomTensor2[VV, EE];
]
First@AbsoluteTiming[
Sum2 = S1 + S2 + S3;
]


0.0698053

0.0164063

In particular the addition operation is 15 times faster. Should be worth to rearrange the data a bit.