如何让numpy函数可直接接受自定义Mytensor类对象作为输入

实现方案 #

你可以通过NumPy原生支持的__array_ufunc__魔术方法实现自定义类和NumPy函数的适配，完整实现代码如下：

import numpy as np

class Mytensor:
    def __init__(self, data):
        self.data = np.array(data)
        self.parents = []
    
    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        # 解析输入参数，拆分Mytensor对象的data属性和普通参数
        processed_inputs = []
        current_parents = []
        for item in inputs:
            if isinstance(item, Mytensor):
                processed_inputs.append(item.data)
                current_parents.append(item)
            else:
                processed_inputs.append(item)
        
        # 执行原生NumPy运算
        calc_result = ufunc(*processed_inputs, **kwargs)

        # 针对不同运算做差异化返回
        if ufunc == np.dot:
            # 点积运算返回Mytensor实例，绑定父节点用于计算图追踪
            new_tensor = Mytensor(calc_result)
            new_tensor.parents = current_parents
            return new_tensor
        # 其余运算直接返回计算结果，适配np.exp等基础调用需求
        return calc_result

效果测试 #

• 基础NumPy函数调用测试

foo = Mytensor(2)
print(np.exp(foo))
# 输出：7.38905609893065，和np.exp(2)结果一致

• 计算图点积运算测试

t1 = Mytensor([1, 2])
t2 = Mytensor([3, 4])
dot_output = np.dot(t1, t2)
print(type(dot_output)) 
# 输出：<class '__main__.Mytensor'>
print(dot_output.data)
# 输出：11
print(len(dot_output.parents))
# 输出：2，parents属性包含t1、t2两个父节点

如果后续需要对更多NumPy函数做自定义返回逻辑，只需要在__array_ufunc__方法中增加对应ufunc的判断分支即可。

内容的提问来源于stack exchange，提问作者Sam-gege