Compile Flatten with list as second argument

Question

If you try to evaluate code in which Flatten uses a List as second argument, like the following:

foo = Compile[{{u, _Real, 2}},
  Flatten[u, {1, 2}]
  ]

Mathematica will complain with the following error message:

Compile::cpint: {1,2} at position 2 of Flatten[u,{1,2}] should be a machine-sized integer; evaluation will use the uncompiled function.

Using an integer as second argument of Flatten indeed gives no error.

Why is this? And is there a way around it?

My particular use case is to use Flatten to convert a TensorProduct structure into a structure similar to that given by KroneckerProduct, that is, to do something like in the following (but for any number of matrices in the TensorProduct structure):

fooMatrix = TensorProduct[
    Array[a, {2, 2}], Array[b, {2, 2}], Array[c, {2, 2}]
  ];
Flatten[fooMatrix, {{1, 3, 5}, {2, 4, 6}}]

It also seems that ArrayFlatten, which could probably be tweaked to work here instead of Flatten, is not compilable (as also shown by trying to compile it and observe the CompilePrint).

Answer 1

First of all I'd like to point out that, even if a function is compilable, its functionality is usually limited inside Compile. It's not rare to find a valid syntax of a compilable function isn't compilable. I think in this case the description after Compile::cpint is clear: the 2nd argument of Flatten can only be an integer inside Compile i.e. the 4th syntax in the document of Flatten is just not supported inside Compile, period.

Then how to circumvent this? The only solution I can think out is to self-implement it with Table, with the technique mentioned here to make it as fast as the built-in:

iter[j_] := {i[j], dim[j]}

lst = {1, 3, 5, 2, 4, 6};

cf = With[{g = Compile`GetElement}, 
     Hold@Compile[{{u, _Real, 6}}, 
           Module[{dim = Dimensions@u}, 
            Partition[Flatten@Table[##], dim[2] dim[4] dim[6]]], CompilationTarget -> C, 
           RuntimeOptions -> "Speed"] &[
        g[u, ## & @@ i /@ Range@6], ## & @@ iter /@ lst] /. 
       i[j_] :> RuleCondition@ToExpression["i" <> ToString@j] /. dim[j_] :> g[dim, j]] //
     ReleaseHold // Last;

test = RandomReal[{0, 1}, {10, 10, 10, 10, 10, 10}];

rst1 = Flatten[test, {{1, 3, 5}, {2, 4, 6}}]; // AbsoluteTiming
(* {0.015607, Null} *)

rst2 = cf@test; // AbsoluteTiming
(* {0.015645, Null} *)

rst2 == rst1
(* True *)

Answer 2

I'm not sure that trying to Compile Flatten is going to buy you any speed up. For instance, let's suppose we were interested in using just Flatten with 1 argument. Using the following data:

data=RandomReal[1,{10,10,10,10,10,10}];

Let's compare a compiled versus non-compiled version:

fc = Compile[{{x, _Real, 6}}, Flatten[x]];
r1 = fc[data]; //AbsoluteTiming
r2 = Flatten[data]; //AbsoluteTiming
r1 === r2

{0.006045, Null}

{0.002228, Null}

True

So, the non-compiled version is faster. Even if a 2-arg version of Flatten were compilable, I think the non-compiled version would probably be faster.