# What's the fastest way to take Mean of a Tensor in given slot?

One should expect that the implementation of the built-in function Mean was quick. However, as I recently observed, this is not true at all, in particular for matrices. Just compare Mean to the following simple function:

mean[A_?TensorQ] := ConstantArray[1./Length[A], Length[A]].A


A test with a decently sized matrix:

A = RandomReal[{-1, 1}, {4000, 5000}];
a = Mean[A]; // AbsoluteTiming//First
b = mean[A]; // AbsoluteTiming//First
Max[Abs[a - b]]

(* 0.220566 *)
(* 0.008163 *)
(* 1.71738*10^-16 *)


Wow, for a built-in function, this is really, really bad. For higher rank matrices, this gap is not as pronounced but it is still there:

A = RandomReal[{-1, 1}, {400, 500, 300}];
a = Mean[A]; // AbsoluteTiming // First
b = mean[A]; // AbsoluteTiming // First
Max[Abs[a - b]]

(* 0.110258 *)
(* 0.024915 *)
(* 2.22045*10^-16 *)


Okay, for the first (and also for the last slot), there is this workaround by using Dot. What about a slot in between? E.g. how to compute Total[A,{2}]/Dimensions[A][] as fast as possible?

PS.: The same mess appears with Total, a function that I usually value quite high, in particular since Total allows for level specs.

• Using exact values in the targets and in the function reverses the performance dramatically. Hypothesis: I'd venture the dot with machine precision takes a path using Intel MKL, while Mean/Total, being built to handle general forms (e.g. symbolic, mixed) dont. – ciao Oct 22 '17 at 8:43
• @ciao You are totally right about exact arithmetic. I really aim for PackedArrays. What I wanted to point out is that it might be a good idea to overload Mean for PackedArrays to something that actually uses Intel MKL. – Henrik Schumacher Oct 22 '17 at 9:15

Here is a compiled version of Total that is faster for intermediate levels:

total[d_, l_] := With[
{o = ConstantArray[1., Dimensions[d][[l]]], n = Length@Dimensions@d-l+1},
Switch[l,
1, o.d,

Length@Dimensions@d, d.o,

_,
With[{fc = Compile[{{data, _Real, n}}, o.data, RuntimeAttributes->{Listable}]},
fc[d]
]
]
]


d = RandomReal[{-1, 1}, {400, 500, 300}];

r1 = total[d, 2]; //RepeatedTiming
r2 = Total[d, {2}]; //RepeatedTiming

r1 == r2


{0.0632, Null}

{0.11, Null}

True

I will leave the modification to compute the mean instead to you.

• Shiny! I wasn't aware that we can control the threading level of listable compiled functions this way. Thank you, this answers also several further questions! =o) – Henrik Schumacher Oct 23 '17 at 7:56

First, don't use TensorQ. It is undocumented and may change in the future. The documented function to use is ArrayQ. Second, here an arbitrary-rank version of your function,

mean[A_?ArrayQ, slot_] :=  With[
{d = Part[Dimensions[A], slot]},
ConstantArray[1./d, d].Transpose[A, 1 <-> slot]
]


Unfortunately, it is slower than Total for higher-rank, because the cost of rearranging memory is too high. I suspect that Carl's approach is probably the best you can do, but I can't immediately prove that.

• Thanks for the hint on ArrayQ. And you are right, transposing is not an option... – Henrik Schumacher Oct 24 '17 at 5:58