[FIXED] How to implement infix operator matrix multiplication in Python 2?

Issue

I have been using python 3.7 for a few months, but I recently had to shift to python 2.7. Since I am developing scientific code, I heavily rely on the use of infix operator @ to multiply nd-arrays. This operator was introduced with python 3.5 (see here), therefore, I cannot use it with my new setup.

The obvious solution is to replace all M1 @ M2 by numpy.matmul(M1, M2), which severely limits my code's readability.

I saw this hack which consists in defining an Infix class allowing to create custom operators by overloading or and ror operators. My question is: How could I use this trick to make an infix |at| operator working just like @?

What I tried is:

import numpy as np

class Infix:
    def __init__(self, function):
        self.function = function
    def __ror__(self, other):      
        return Infix(lambda x, self=self, other=other: self.function(other, x))
    def __or__(self, other):
        return self.function(other)
    def __call__(self, value1, value2):
        return self.function(value1, value2)

# Matrix multiplication
at = Infix(lambda x,y: np.matmul(x,y))

M1 = np.ones((2,3))
M2 = np.ones((3,4))

print(M1 |at| M2)

When I execute this code, I get :

ValueError: operands could not be broadcast together with shapes (2,3) (3,4) 

I think I have an idea of what is not working. When I only look at M1|at, I can see that it is a 2*3 array of functions:

array([[<__main__.Infix object at 0x7faa1c0d6da0>,
        <__main__.Infix object at 0x7faa1c0d6860>,
        <__main__.Infix object at 0x7faa1c0d6828>],
       [<__main__.Infix object at 0x7faa1c0d6f60>,
        <__main__.Infix object at 0x7faa1c0d61d0>,
        <__main__.Infix object at 0x7faa1c0d64e0>]], dtype=object)

This is not what I expected, since I would like my code to consider this 2d-array as a whole, and not element-wise...

Does anybody have a clue of what I should do?

PS: I also considered using this answer, but I have to avoid the use of external modules.

Solution

I found the solution to my problem here.

As suggested in the comments, the ideal fix would either be to use Python 3.x or to use numpy.matmul, but this code seems to work, and even has the right precedence :

import numpy as np

class Infix(np.ndarray):
    def __new__(cls, function):
        obj = np.ndarray.__new__(cls, 0)
        obj.function = function
        return obj
    def __array_finalize__(self, obj):
        if obj is None: return
        self.function = getattr(obj, 'function', None)
    def __rmul__(self, other):
        return Infix(lambda x, self=self, other=other: self.function(other, x))
    def __mul__(self, other):
        return self.function(other)
    def __call__(self, value1, value2):
        return self.function(value1, value2)

at = Infix(np.matmul)

M1 = np.ones((2,3))
M2 = np.ones((3,4))
M3 = np.ones((2,4))

print(M1 *at* M2)
print(M3 + M1 *at* M2)

Answered By - Nicolas