FunFact icon indicating copy to clipboard operation
FunFact copied to clipboard

Support Ellipsis in indices?

Open yhtang opened this issue 3 years ago • 3 comments
trafficstars

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]

yhtang avatar Dec 30 '21 18:12 yhtang

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.

yhtang avatar Dec 30 '21 18:12 yhtang

Continued the discussion here @yhtang

campsd avatar Feb 18 '22 14:02 campsd

Per discussion with @campsd, it might be better to collect more usage scenarios before deciding on to what extent should this feature be supported.

yhtang avatar Mar 17 '22 22:03 yhtang