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).

Where is the form Flatten[u, {1, 2}] documented? I cannot find it. — Marius Ladegård Meyer, Mar 14 '17 at 20:30
@MariusLadegårdMeyer interestingly, I just noticed it isn't, not directly at least. It seems to work the same as Flatten[u, {{1, 2}}]. Using the latter form gives the same problem with Compile though (and is indeed the one I'm interested in using, I used the one with a rank-1 list as second argument here just to make a simple reproducible example) — glS, Mar 14 '17 at 20:37

score 7 · Accepted Answer · edited Apr 13 '17 at 12:55

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 *)

score 4 · Answer 2 · answered Mar 15 '17 at 00:14

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.

While I see your point, if you want to compile a more complex function that uses Flatten as part of the computation, this prevents the whole function from being compiled efficiently. This actually arose while trying to compile the partialTrace function of my other question — glS, Mar 15 '17 at 00:21

Compile Flatten with list as second argument

2 Answers2