如何将对象传递给 numpy 点函数

假设我已经定义了我的对象：

import numpy as np
class myTensor:
    def __init__(self,data):
        self.data=np.array(data)
        self.parent=[]

如何将“myTensor”作为输入传递给 np.dot？例如：

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

现在，如果我打电话

t3=np.dot(t1,t2)

我期望t3是另一个myTensor对象，其t3.data是数组[1,2]和[3,4]的点积，即11； t3.parents 是包含 t1 和 t2 的列表。

感谢您的任何建议！

查看有关如何编写自定义数组容器的教程：https://numpy.org/devdocs/user/basics.dispatch.html 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

您现在可以运行 np.exp 或 np.dot 并返回一个 Tensor 对象：

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