einx.torch.adapt_numpylike_elementwise

einx.torch.adapt_numpylike_elementwise(op)

Adapts an operation with a Numpy-like element-wise signature to einx notation.

The operation is expected to have the following Numpy-like signature:

def op(*tensors: Tensor) -> Tensor:
    ...

All input tensors are guaranteed to be broadcastable according to Numpy’s broadcasting rules. They are additionally guaranteed to have the same number of dimensions. The function should return a tensor with the broadcasted shape of the input tensors.

Parameters:

op – The operation that will be adapted to einx notation.

Returns:

A new operation that follows einx notation and internally invokes the original operation.