Fastest way to ArrayFlatten a 2D array of 2D arrays, both for Sparse and Dense

Question

I am using Mathematica 12 and want to know what is the fastest way to transform a 2D array of 2D arrays both for dense and sparse arrays. I remember looking for this a few versions of Mathematica ago, when ArrayFlatten was not that good. Since then I have been using option one below (which I found somewhere here, but I can't find the question anymore - it had other suggestions as well). But a quick check shows that for SparseArrays that is not the case anymore. Is there anything better than ArrayFlatten

Also, why is the SparseArray version of 3 not the same as the other two, unless I make them dense?

dim = 5; eles = 50;
sparse = Table[
   KroneckerProduct[RandomReal[{-10, 10}, {eles, eles}], 
    IdentityMatrix[eles, SparseArray]], {ii, 1, dim}, {jj, 1, dim}];
Dimensions[sparse]

{5, 5, 2500, 2500}

RepeatedTiming[
 jSparse1 = 
   Apply[Join[##, 2] &, 
    Table[Join @@ sparse[[All, ii]], {ii, 1, dim}]];]

{0.04, Null}

RepeatedTiming[jSparse2 = ArrayFlatten[sparse];]

{0.0094, Null}

RepeatedTiming[
 jSparse3 = 
   SparseArray`SparseBlockMatrix[{{i_, j_} :> sparse[[i, j]]}, {dim, 
     dim}];]

{0.23, Null}

Dimensions /@ {jSparse1, jSparse2, jSparse3}

{{12500, 12500}, {12500, 12500}, {12500, 12500}}

jSparse1 === jSparse2

True

jSparse1 === jSparse3

False

jSparse2 === jSparse3

False

Normal[jSparse1] === Normal[jSparse2] === Normal[jSparse3]

True

And here is the dense version with same conclusion

dim = 5; eles = 50;
dense = 
  Table[KroneckerProduct[RandomReal[{-10, 10}, {eles, eles}], 
    IdentityMatrix[eles]], {ii, 1, dim}, {jj, 1, dim}];
Dimensions[dense]

{5, 5, 2500, 2500}

RepeatedTiming[
 jDense1 = 
   Apply[Join[##, 2] &, 
    Table[Join @@ dense[[All, ii]], {ii, 1, dim}]];]

{1.251, Null}

In[5]:= RepeatedTiming[jDense2 = ArrayFlatten[dense];]

{0.61, Null}

RepeatedTiming[
 jDense3 = 
   Normal[SparseArray`SparseBlockMatrix[{{i_, j_} :> 
       dense[[i, j]]}, {dim, dim}]];]

{2.3, Null}

Dimensions /@ {jDense1, jDense2, jDense3}

{{12500, 12500}, {12500, 12500}, {12500, 12500}}

jDense1 === jDense2

True

jDense1 === jDense3

False

jDense2 === jDense3

False

Max[jDense1 - jDense2]

0.

Max[jDense1 - jDense3]

0.

Max[jDense2 - jDense3]

0.

and again the third version produces a result that apparently is not identical with the other two, though numerically they seem to be the same.

Answer 1

Applying SparseArray`SparseArraySort to jSparse3 will make it the same as jSparse1 and jSparse2. Also notice that jSparse1 == jSparse3 returns True.

SameQ (===) compares the internal structure of the SparseArrays, which are the "ColumnIndices", "RowPointers", "NonzeroValues", and "Background" properties. At least, that's the defining properties for vectors and matrices. (I am not totally sure about higher order tensors.)

The problem is that not every Mathematica SparseArray is stored in the standard compressed row storage (CRS) format. But SparseArray`SparseArraySort converts to the standard format. The standard requires the column indices in each row to be ascending; jSparse3 violates this:

row = RandomInteger[Length[jSparse3]];
rp1 = jSparse3["RowPointers"][[row]];
rp2 = jSparse3["RowPointers"][[row + 1]];
OrderedQ@jSparse3["ColumnIndices"][[rp1 + 1 ;; rp2]]

False

Sorting the column indices is not for free, that's why Mathematica sometimes suppresses the ordering. But that makes it also hard to work with such sparse arrays, a further reason why SparseArray`SparseBlockMatrix should be avoided.