einx.mlx.adapt_numpylike_reduce

einx.mlx.adapt_numpylike_reduce(op)

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

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

def op(tensor: Tensor, axis: Tuple[int]) -> Tensor:
    ...

It should return a tensor with the same shape as the input tensor, but with axes at positions specified in axis removed.

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.