Implementation of the Einsum operator

[MarcusGrass]

Hi I'm trying to run an onnx-model that uses the Einsum operator as described here https://github.com/onnx/onnx/blob/main/docs/Operators.md#Einsum added to onnx v12, which is currently unimplemented here in Tract, is implementing that operator in progress?

I looked into it and without much knowledge of how operators function, it seems out of my reach. I looked at the CumSum operator https://github.com/sonos/tract/blob/main/onnx/src/ops/cumsum.rs an operation that is fairly easy to explain and understand, and can be implemented in about a 100 lines, and that might as well have been hieroglyphs to me.

Einsum is from my reading the above onnx docs, fairly complex to implement. Looking at the above (onnx link) description it requires a string as an input describing the equation. The string is requirements are relaxed as well which makes it tougher to parse, ie. "if 'n' is left out, it means that implicitly...". Aside from that it takes and arbitrary amount of input tensors and returns an output tensor.

The actual mathematical operations converted to onnx-language is my major holdup here while parsing the input string is absolutely doable. If there's someone here with the knowledge of how to implement those I could definitely take a shot at the parsing.

[kali]

Hello, thanks for your interest in tract!

There is no effort to implement Einsum at this stage from our side, you are more than welcome to give it a shot. My understanding is that Einsum equations can always be reinterpreted in terms of other operators, like matrix product and axes permutation (aka MatMul and Transpose). This is good news, because you should not have to worry too much about performance.

For operators like these, you should probably be looking at implementations of the Expansion trait in tract-onnx. CumSum is really not an easy one, I won't get into it. Better to have a look at: https://github.com/sonos/tract/blob/play-with-spleeter/onnx/src/ops/nn/batch_norm.rs https://github.com/sonos/tract/blob/play-with-spleeter/onnx/src/ops/array/squeeze.rs

Basically implementing Expansion means two things: 1/ provide rules that capture the relation between inputs and output shapes as closely as practical (rust API has not aged that well, but there are lots of examples. I can help). 2/ implement the translation of the actual operator logic using tract-core operators in the wire method.

With me so far?

[kali]

#731 provides support for Einsum. Optimisation still pending, so better not try to use it as the an general purpose affine operator at this stage.

[kali]

We do have a reasonable set of optimisation in place for binary Einsum at this stage. Not for the other cases, but they are not so frequent, we will address them when/if we see evidence of usage.