How to pass objects to numpy dot function

user_16876370 0
Viewed 117

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' as inputs to np.dot? For example:

t1=myTensor([1,2])
t2=myTensor([3,4])

Now, if I call

t3=np.dot(t1,t2)

I expect t3 to be another myTensor object, whose t3.data is the dot product of array [1,2] and [3,4], that is, 11; and t3.parents to be the list containing t1 and t2.

Thanks for any suggestions!!

numpy python
1 Answers

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}')
user_142200870
Related