FunFact
FunFact copied to clipboard
Support Ellipsis in indices?
Is your feature request related to a problem? Please describe.
Right now, users have to explicitly provide all the indices to perform Einstein operations. In cases where simple broadcasting is obvious, it might be more convenient to allow ellipsis indexing. numpy.einsum already supports this in the einspec strings.
Describe the solution you'd like
a = ff.tensor('a', 3, 2, 3, 4)
b = ff.tensor('b', 3, 2, 4, 6)
# current
i, j, k, r, s = ff.indices('i, j, k, r, s')
a[~i, ~j, k, r] * b[~i, ~j, r, s]
# proposed
i, j, k = ff.indices('i, j, k')
a[..., i, j] * b[..., j, k]
A possible approach to implement it would be to introduce a special value for the indices AST primitive, and then rewrite it using lhs&rhs information in the __call__ of the new index propagator.
Continued the discussion here @yhtang
Per discussion with @campsd, it might be better to collect more usage scenarios before deciding on to what extent should this feature be supported.