Tensor contraction?


Dear all,

is there a (simple) way to do tensor contraction with NDArray? Say I have a tensor T^{ij}_{kl} (i.e. a 4D NDArray with dim(j) = dim(k) ), I want to perform contraction over the j,k indices. For example, if dim(j) == dim(k) = 5 I need to perform the operation:

Result^i_k = \sum_{j=0}^{4} T^{ij}_{jk}

this is different than writing nd.sum(T,axis=[2,3]) which sums also for indices j != k.

As another example, for a 2D NDArray

data = nd.ones(shape = [5,5])

temp = 0.0
for i in range(5):
    temp += data[i,i]

here the variable temp (contraction) has the value 5 (correct), while nd.sum(data,axis=[0,1]) results 25 (not contraction).

Thank you very much.