How To Pass Objects To Numpy Dot Function

Suppose I have defined my object: import numpy as np class myTensor: def __init__(self,data): self.data=np.array(data) self.parent=[] How can I pass 'myTensor'

Solution 1:

check the tutorial on how to write custom array container: https://numpy.org/devdocs/user/basics.dispatch.html

from numbers import Number

import numpy as np

HANDLED_FUNCTIONS = {}


class Mytensor():
    def __init__(self, data):
        self.data = np.array(data)
        self.parents = []

    def __repr__(self):
        return f'Tensor: {self.data}'

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        if method == '__call__':
            scalars = []
            for input in inputs:
                if isinstance(input, Number):
                    scalars.append(input)
                elif isinstance(input, self.__class__):
                    scalars.append(input.data)
                else:
                    return NotImplemented
            output = self.__class__(ufunc(*scalars, **kwargs))
            output.parents.extend(scalars)
            return output
        else:
            return NotImplemented

    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        if not all(issubclass(t, self.__class__) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)


def implements(np_function):
    def decorator(func):
        HANDLED_FUNCTIONS[np_function] = func
        return func

    return decorator


@implements(np.dot)
def dot(t1: Mytensor, t2: Mytensor):
    output = Mytensor(np.dot(t1.data, t2.data))
    output.parents.extend([t1, t2])
    return output

You can now run np.exp or np.dot and returns an Tensor object:

if __name__ == '__main__':
    t1 = Mytensor([[1, 2], [3, 4]])
    t2 = Mytensor([[5, 6], [7, 8]])
    output1 = np.exp(t1)
    print(f'result: {output1}')
    print(f'parents: {output1.parents}')

    output2 = np.dot(t1, t2)
    print(f'result: {output2}')
    print(f'parents: {output2.parents}')