Generalize AstSum0S and AstSum and also AstDot0S and AstDot1InS

[Mikolaj]

These would probably take a list of dimensions to sum and to leave alone, maybe in the form of a StaticShX value or similar. An alternative is to always sum first n dimensions (where n comes from SNat @n) and use transposition to recover the more general functionality.

A related task is to then vectorize the generalized two constructors. Once they can be vectorized, we can use them and keep them early, instead of, as we do now, expressing them in terms of simpler operations and recognizing them back again in the contraction phase.

Another related task is to similarly generalize the dual operations related to replication so that we can differentiate the new operations efficiently and readably.

[tomsmeding]

ox-arrays and Futhark handle transpose->sum just fine; ox-arrays does not handle transpose->reshape->sum without renormalisation if the reshaped dimensions are not continuous in the underlying array. So an n-dimensional sum should be sufficient.

Do you have a strong reason for wanting more generalised sums than that?

[Mikolaj]

More generalized than they are now plus transpositions? Yes, my reason is that vectorization of

  AstSum0S :: AstTensor ms s (TKS2 sh x)
           -> AstTensor ms s (TKS2 '[] x)

and probably of

  AstDot0S :: GoodScalar r
           => AstTensor ms s (TKS sh r) -> AstTensor ms s (TKS sh r)
           -> AstTensor ms s (TKS '[] r)

requires an operation that sums only given dimensions. Currently I only have operations that sum the first or all dimensions. Transpositions make it "one or all dimensions", which is not enough.

[tomsmeding]

I see, I forgot that we don't have n-dimensional sums yet. I guess that means that I think n-dimensional sums (sum the outer/inner n dimensions) are sufficient, because combined with transposition, that produces your fully general sum operation.

[Mikolaj]

Yes, indeed. I think even ox-arrays doesn't yet have n-dimensional sums, even though it has n-dimensional replicates. Horde-ad doesn't have either, but needs both.