Using IndexSetVar defined at runtime

[smr97]

Hi everyone!
I have a use-case where I want to "slice" into a high dimensional tensor at specific points in the leading dimension.
However, these points change at run-time. I see from the code that IndexSetVar is formed using an integer given to the IndexVar before the TACO code is generated (before calling .compile()). Is there a way to compile the einsum expression with a generic slice position, and run the contraction with a different value of the slice position?
Here is an example:

Tensor<double> threedee = Tensor<double>("3d", {X, Y, Z});
Tensor<double> matr = Tensor<double>("mat", {Z, A});
Tensor<double> res = Tensor<double>("res", {Y, A});
for(int x = 0; x < X; x++){
    res(y, a) = threedee(x_var({x}), y, z) * res(y, a);
    res.evaluate();
}

This is how I currently do it, but it is slow, because the TACO compilation happens for every iteration of the for loop.
I want the code to be generated just once, with an argument x, and I can pass in a different value of x each iteration of the loop. So, something like this:

Tensor<double> threedee = Tensor<double>("3d", {X, Y, Z});
Tensor<double> matr = Tensor<double>("mat", {Z, A});
Tensor<double> res = Tensor<double>("res", {Y, A});
int x = 0;
res(y, a) = threedee(x_var({x}), y, z) * res(y, a);
res.compile();
for(int iter = 0; iter < X; iter++){
    res.assemble(x = iter);
    res.compute(x = iter);
}

Maybe this is already implemented in a different way? I am not sure what's the fastest way to run this, so any help would be appreciated!!