Supporting same indices in einsum operate or evaluate()

[ByzanTine]

Looking at the einsum documentation warning, you suggests that "the same subscript cannot appear more than once in a given operand". So I tried in pt.evaluate() as well, but it seems there is no luck as well. The use case being einsum that involves "...kk->k.." for example, or simpler things like trace.

import numpy as np
import pytaco as pt
t = pt.tensor([5, 5], pt.csr)
pt.evaluate("S = A(i, i)", t).to_array()

I got errors like

/tmp/taco_tmp_u95yF1/neq7js40dzrq.c: In function ‘compute’:
/tmp/taco_tmp_u95yF1/neq7js40dzrq.c:126:30: error: ‘iA00’ undeclared (first use in this function); did you mean ‘iA0’?
   for (int32_t iA0 = A02_pos[iA00]; iA0 < A02_pos[(iA00 + 1)]; iA0++) {
                              ^~~~
                              iA0
/tmp/taco_tmp_u95yF1/neq7js40dzrq.c:126:30: note: each undeclared identifier is reported only once for each function it appears in
Traceback (most recent call last):
  File "test.py", line 4, in <module>
    pt.evaluate("S = A(i, i)", t).to_array()
  File "/home/jiayu/workspace/taco/build/lib/pytaco/pytensor/taco_tensor.py", line 403, in to_array
    return to_array(self)
  File "/home/jiayu/workspace/taco/build/lib/pytaco/pytensor/taco_tensor.py", line 699, in to_array
    return np.array(t.to_dense(), copy=True)
  File "/home/jiayu/workspace/taco/build/lib/pytaco/pytensor/taco_tensor.py", line 396, in to_dense
    new_t[idx_vars] = self[idx_vars]
  File "/home/jiayu/workspace/taco/build/lib/pytaco/pytensor/taco_tensor.py", line 317, in __getitem__
    return self._tensor[index]
RuntimeError: Error at /home/jiayu/workspace/taco/src/codegen/module.cpp:154 in compile:
 Compilation command failed:
cc -O3 -ffast-math -std=c99 -shared -fPIC /tmp/taco_tmp_u95yF1/neq7js40dzrq.c  -o /tmp/taco_tmp_u95yF1/neq7js40dzrq.so -lm
returned 256

Am I writing the evaluate expressions wrong or there is no support for trace like operation at all? And if latter, wondering how complicated would it be to implement such supports?

[rawnhenry]

If I remember correctly, the einsum parser actually just generates index notation. I believe index notation in general does not support using the same index variable to access a tensor meaning we can't do things like traces currently.

I don't recall why that is the case so I'm not sure how complicated it would be to add support for that. Perhaps the docs for pt.evaluate should be updated with a similar warning.

[ByzanTine]

Thanks for the response, then I am wondering how hard would it be for index notation to support using the same index var to access a tensor of different dimensions.

[rawnhenry]

I think the main issue is having the compiler generate the correct code to locate in accesses that are indexed with repeated vars. (For dense its easy but for ordered sparse you have to binary search) The bare bones way to do this is to probably change the lowerer (and relax whatever assertions fail before then). However, this way may not be the cleanest. I think something like this is nontrivial to implement but definitely doable.