CMU 10-414/714: Deep Learning Systems --hw1

hw1

needle库中的两个重要文件

1. autograd.py文件：实现计算图框架、自动微分框架（automatic differentiation framework）
2. ops.py文件：实现各种算子。反向传播的粒度很细，会以一个个算子为单位，而不是以网络层为单位。其实现基础是numpy，后面会自己实现自己的运算库

needle中的一些核心类

1. Value：计算图中的节点，代表计算过程中的数值
2. Op：计算图的算子，即加减乘除等操作。其通过compute()函数进行前向计算，通过gradient()函数进行反向传播
3. Tensor：是Value的子类，表示计算图中的一个多维数组。课程提供了一些功能（如操作符重载）来保证用tensor进行数学运算
4. TensorOp：是Op的子类，表示返回Tensor的算子

实现功能

1. 实现前向传播和反向传播
   • compute()函数实现前向运算，运算的输入都是numpy.ndarray（后面的实验会实现自己的NDArray，因此这里为了对齐用import numpy as array_api），即compute()函数只处理raw data，而不处理自动求导的Tensor
   • gradient()函数实现反向传播，其输入为Tensor，且函数内的运算都是TensorOp算子
   • 需要实现运算的算子如下。这里代码其实没有什么高深的，只要搞清楚每个算子是怎么求导的就好了，比如PowerScalar算子实现$a^{scalar}$，其gradient()函数就是返回$(out\_grad\times self.scalar\times lh{s}^{self.scalar-1})$
     • PowerScalar：乘方
     • EWiseDiv：除法
     • DivScalar：张量除以标量
     • MatMul：矩阵乘法
       • 梯度必须和输入数据的维度一致：损失函数L对输入数据X的梯度，告诉我们在每个特征维度上，改变输入数据时损失函数的变化情况。若想要计算L对每个输入特征的梯度，这些梯度将组成一个与输入数据X相同形状的张量，每个元素表示在相应特征维度上的变化率
     • Summation：在指定的维度上求和（在哪个维度上求和，导数就在哪个维度上扩展成与输入张量相同的形状）
     • BroadcastTo：broadcast成指定的形状（在哪个维度上扩展，导数就在哪个维度上求和后reshape成与输入张量相同的形状）
     • Reshape：转变成之后定的形状（导数是reshape成输入张量相同的形状）
     • Negate：取负数
     • Transpose：交换两个维度（导数仍是交换这两个维度）
2. 拓扑排序：原理很简单，数据结构里学过。就是深度优先遍历node_list，如果一个node没有被visited，就深度优先遍历它的inputs结点，直到一个入度为0的结点，就将其加入topo_order中

def find_topo_sort(node_list: List[Value]) -> List[Value]:
    visited = set()
    topo_order = []
    for node in node_list:
      if node not in visited:
        topo_sort_dfs(node, visited, topo_order)
    return topo_order

def topo_sort_dfs(node, visited, topo_order):
    if node in visited:
      return

    for next in node.inputs:
      topo_sort_dfs(next, visited, topo_order)
    
    visited.add(node)
    topo_order.append(node)

3. 实现reverse mode differentiation：

def compute_gradient_of_variables(output_tensor, out_grad):
  # 创建一个空字典，用于存储每个节点在计算图中的梯度信息
  node_to_output_grads_list: Dict[Tensor, List[Tensor]] = {}
  # 将输出张量的梯度存储在字典中，作为输出节点的梯度
  node_to_output_grads_list[output_tensor] = [out_grad]

  # find_topo_sort([output_tensor])从输出张量开始计算图的拓扑排序（因为输出张量通过node.inputs连着整张图，所以可以实现）
  reverse_topo_order = list(reversed(find_topo_sort([output_tensor])))

  for node in reverse_topo_order:
    # 计算当前节点的梯度：就是把所有偏导数加到一起
    sum_grad = node_to_output_grads_list[node][0]
    for t in node_to_output_grads_list[node][1:]:
      sum_grad = sum_grad + t 
    node.grad = sum_grad
	
	# 叶子节点即没有op的节点，跳过后面的for循环
    if node.is_leaf():
      continue

	# 计算节点的输入节点相对于
    for i, grad in enumerate(node.op.gradient_as_tuple(node.grad, node)):
      input_ = node.inputs[i]
      if input_ not in node_to_output_grads_list:
        node_to_output_grads_list[input_] = []
      node_to_output_grads_list[input_].append(grad)

举个例子有助于理解自动微分的过程原理：

4. 重新实现softmax loss：使用Tensor且y是one-hot向量组成的$m\times k$矩阵

def softmax_loss(Z, y_one_hot):
    a = ndl.ops.summation(Z*y_one_hot)
    b = ndl.ops.summation(ndl.ops.log(ndl.ops.summation(ndl.ops.exp(Z), axes=(1,))))
    return (b-a) / Z.shape[0]

5. 重新为一个两次神经网络实现SGD
   与hw0的区别：W1和W2变成Tensor，X和y还是numpy array；hw0中是手动计算反向传播公式然后用代码实现，这次则是使用自动微分机制，只需算出loss值，然后调用.backward()即可自动对每个参数求导，再对W1和W2更新即可

def nn_epoch(X, y, W1, W2, lr = 0.1, batch=100):
    for i in range(X.shape[0]//batch+1):
      # 首先不变的，采样batch
      start, end = i * batch, min((i+1)*batch, X.shape[0])
      m = end - start
      if m == 0:
          break
      Xb, yb = X[start:end], y[start:end]

      # hw0这里是计算梯度，是通过手动计算得出数学公式然后实现的
      # 现在则是使用自动微分机制，算出loss后直接.backward()即可
      # 下面并没有对X和y设置自动微分，因为参数更新过程中只需要对参数矩阵W1和W2自动微分即可，不需要对输入数据和得出来的结果微分（node_to_grad第一个值是loss）
      Xb, yb = X[start:end], y[start:end]
      Xb = ndl.Tensor(Xb, requires_grad=False)
      Z1 = Xb @ W1
      A1 = ndl.ops.relu(Z1)
      Z2 = A1 @ W2
      y_one_hot = np.zeros(Z2.shape, dtype=np.float32)
      y_one_hot[np.arange(Z2.shape[0]), yb] = 1
      y_one_hot = ndl.Tensor(y_one_hot, requires_grad=False)
      loss = softmax_loss(Z2, y_one_hot)
      loss.backward()
      
      # 更新参数矩阵
      W1.data = W1.data - lr * ndl.Tensor(W1.grad.numpy().astype(np.float32)) 
      W2.data = W2.data - lr * ndl.Tensor(W2.grad.numpy().astype(np.float32))
    return W1, W2

相关知识补充

一、pytorch中的正向传播、反向传播和计算图

1. Tensor
   pytorch提供的autograd包能够根据输入和前向传播过程自动构建计算图，并执行反向传播。
   
   Tensor是autograd包的核心类，若将其属性.requires_grad设为True，它将开始追踪在其上的所有操作（这样就可以利用链式法则进行梯度传播了）。完成计算后，可调用.backward()来完成所有的梯度计算。此Tensor的梯度将累积到.grad属性中。
   
   Function是另一个核心类，Tensor和Function结合就可以构建一个记录有整个计算过程的有向无环图。每个Tensor的.grad_fn属性记录着这个Tensor最近是通过哪一个运算得到的，若其不是通过某个运算得到的，.grad_fn属性值为None，像grad_fn属性值为None的称为叶子节点

x = torch.ones(2, 2, requires_grad=True)
y = x + 2
z = y * y * 3
out = z.mean()
print(z, out)

输出为：

tensor([[27., 27.],
        [27., 27.]], grad_fn=<MulBackward>) tensor(27., grad_fn=<MeanBackward1>)

2. 梯度
   
   $$out=\frac{1}{4}\sum _{i=1}^4{z}_i=\frac{1}{4}\sum _{i=1}^{4}3(x_i+2{)}^2$$ 所以
   $$\frac{\partial out}{\partial x_i}{|}_{x_i=1}=\frac{9}{2}=4.5$$
   
   数学上，若有一个函数值和自变量都为向量的函数$\vec{y}=f(\vec{x})$，那么$\vec{y}$关于$\vec{x}$的梯度就是一个雅可比矩阵：
   $$J=\left(\begin{array}{ccc}\frac{\partial y_1}{\partial x_1}& \cdots & \frac{\partial y_1}{\partial x_n}\\ ⋮& \ddots & ⋮\\ \frac{\partial y_m}{\partial x_1}& \cdots & \frac{\partial y_m}{\partial x_n}\end{array}\right)$$
   而torch.autograd这个包就是用来计算一些雅可比矩阵的乘积的。如$v$是一个标量函数$l=g(\vec{y})$的梯度：
   $$v=\left(\begin{array}{ccc}\frac{\partial l}{\partial y_1}& \cdots & \frac{\partial l}{\partial y_m}\end{array}\right)$$ 那么根据链式法则我们有关于的雅克比矩阵就为: $$vJ=\left(\begin{array}{ccc}\frac{\partial l}{\partial y_1}& \cdots & \frac{\partial l}{\partial y_m}\end{array}\right)\left(\begin{array}{ccc}\frac{\partial y_1}{\partial x_1}& \cdots & \frac{\partial y_1}{\partial x_n}\\ ⋮& \ddots & ⋮\\ \frac{\partial y_m}{\partial x_1}& \cdots & \frac{\partial y_m}{\partial x_n}\end{array}\right)=\left(\begin{array}{ccc}\frac{\partial l}{\partial x_1}& \cdots & \frac{\partial l}{\partial x_n}\end{array}\right)$$
   
   注意：grad在反向传播过程中是累加的，即在某一次前向传播和反向传播过程中计算得到了某些参数的梯度值，然后在下一次前向传播和反向传播中，会再次计算这些参数的梯度值。此时新计算得到的梯度值会将与之前的梯度值相加，而不是覆盖掉之前的梯度值。
   
   这个功能的好处是可以实现低显存跑大batchsize，详细可见此篇文章
3. 正向传播计算图
   假设神经网络如下：
   样本为$\mathbf{x}\in {\mathbb{R}}^d$，中间变量为$\mathbf{z}={\mathbf{W}}^{(1)}\mathbf{x}$，其中${\mathbf{W}}^{(1)}\in {\mathbb{R}}^{h\times d}$是隐藏层的权重参数。
   将中间变量$\mathbf{z}\in {\mathbb{R}}^h$通过激活函数$\varphi$后，得到隐藏激活向量$\mathbf{h}=\varphi (\mathbf{z})$。隐藏层变量$\mathbf{h}$也是一个中间变量。
   假设输出层参数只有权重${\mathbf{W}}^{(2)}\in {\mathbb{R}}^{q\times h}$，可得到向量长度为$q$的输出层变量$\mathbf{o}={\mathbf{W}}^{(2)}\mathbf{h}$.
   假设损失函数为$\ell$，且样本标签为$y$，可以计算出单个数据样本的损失项$\mathbf{L}=\ell (\mathbf{o},y)$.
   根据$L_2$范数正则化的定义，给定超参数$\lambda$，正则化项即$s=\frac{\lambda }{2}\left(|{\mathbf{W}}^{(1)}{|}_F^2+|{\mathbf{W}}^{(2)}{|}_F^2\right)$.
   最终模型在给定的数据有样本上带正则化的损失为$J=L+s$，$J$即为给定数据样本的目标函数
   
   通过绘制计算图来可视化算子和变量在计算中的依赖关系，如下图。正方形表示变量、圆圈表示操作符。左下角表示输入、右上角表示输出
   
4. 反向传播
   反向传播就是计算神经网络参数梯度的方法。总的来说，反向传播依据链式法则，沿着输出层到输入层的顺序，依次计算并存储目标函数有关神经网络各层的中间变量以及参数的梯度。对输入或输出$\mathsf{X},\mathsf{Y},\mathsf{Z}$为任意形状张量的函数$\mathsf{Y}=f(\mathsf{X})$和$\mathsf{Z}=g(\mathsf{Y})$，通过链式法则，有$\frac{\partial \mathbb{Z}}{\partial \mathbf{X}}=\mathrm{prod}\left(\frac{\partial \mathbb{Z}}{\partial \mathbf{Y}},\frac{\partial \mathbf{Y}}{\partial \mathbf{X}}\right)$. 其中$\text{prod}$运算符将根据两个输入的形状，在必要的操作（如转置和互换输入位置）后对两个输入做乘法.
   
   回顾一下本节中样例模型，它的参数是${\mathbf{W}}^{(1)}$和${\mathbf{W}}^{(2)}$，因此反向传播的目标是计算$\frac{\partial J}{\partial {\mathbf{W}}^{(1)}}$和$\frac{\partial J}{\partial {\mathbf{W}}^{(2)}}$。应用链式法则依次计算各中间变量和参数的梯度，其计算次序与前向传播中相应中间变量的计算次序恰恰相反：
   
   首先，分别计算目标函数$J=L+s$有关损失项$L$和正则项$s$的梯度$\frac{\partial J}{\partial L}=1,\frac{\partial J}{\partial s}=1.$
   然后根据链式法则计算目标函数有关输出层变量的梯度$\frac{\partial J}{\partial \mathbf{o}}=\mathrm{prod}\left(\frac{\partial J}{\partial L},\frac{\partial L}{\partial \mathbf{o}}\right)=\frac{\partial L}{\partial \mathbf{o}}.$
   然后计算正则项关于两个参数的梯度：$\frac{\partial s}{\partial {\mathbf{W}}^{(1)}}=\lambda {\mathbf{W}}^{(1)},\frac{\partial s}{\partial {\mathbf{W}}^{(2)}}=\lambda {\mathbf{W}}^{(2)}.$
   然后计算$W^{(2)}$参数的梯度$\frac{\partial J}{\partial {\mathbf{W}}^{(2)}}=\mathrm{prod}\left(\frac{\partial J}{\partial \mathbf{o}},\frac{\partial \mathbf{o}}{\partial {\mathbf{W}}^{(2)}}\right)+\mathrm{prod}\left(\frac{\partial J}{\partial s},\frac{\partial s}{\partial {\mathbf{W}}^{(2)}}\right)=\frac{\partial J}{\partial \mathbf{o}}{\mathbf{h}}^{\top }+\lambda {\mathbf{W}}^{(2)}.$
   沿着输出层向隐藏层继续反向传播，有$\frac{\partial J}{\partial \mathbf{h}}=\mathrm{prod}\left(\frac{\partial J}{\partial \mathbf{o}},\frac{\partial \mathbf{o}}{\partial \mathbf{h}}\right)={\mathbf{W}}^{(2{)}^{\top }}\frac{\partial J}{\partial \mathbf{o}}.$
   由于激活函数$\varphi$是按元素运算的，中间变量$\mathbf{z}$的梯度$\partial J/\partial \mathbf{z}\in {\mathbb{R}}^h$的计算需要使用按元素乘法符$\odot$：$\frac{\partial J}{\partial z}=\mathrm{prod}\left(\frac{\partial J}{\partial \mathbf{h}},\frac{\partial \mathbf{h}}{\partial \mathbf{z}}\right)=\frac{\partial J}{\partial \mathbf{h}}\odot {\varphi }^{\prime }\left(\mathbf{z}\right).$
   最终可得到第一个参数$W^{(1)}$的梯度$\frac{\partial J}{\partial {\mathbf{W}}^{(1)}}=\mathrm{prod}\left(\frac{\partial J}{\partial \mathbf{z}},\frac{\partial \mathbf{z}}{\partial {\mathbf{W}}^{(1)}}\right)+\mathrm{prod}\left(\frac{\partial J}{\partial s},\frac{\partial s}{\partial {\mathbf{W}}^{(1)}}\right)=\frac{\partial J}{\partial \mathbf{z}}{\mathbf{x}}^{\top }+\lambda {\mathbf{W}}^{(1)}$

二、autograd代码解析

各种类

1. Op类：是TensorOp的父类，call、compute、gradient都在子类里各自实现

class Op:
    """Operator definition."""

    def __call__(self, *args):
        raise NotImplementedError()

    def compute(self, *args: Tuple[NDArray]):
        raise NotImplementedError()

    def gradient(self, out_grad: "Value", node: "Value") -> Union["Value", Tuple["Value"]]:
        raise NotImplementedError()

	# 此函数的作用是为了确保 gradient 方法的返回值始终以元组的形式返回。
	# 具体来说，如果 gradient 方法的返回值已经是元组，则直接返回该元组；如果返回值是列表，则将其转换为元组后返回；否则，将其包装成单元素的元组再返回。
	# 这个方法在某种程度上是一种封装，它保证了 gradient 方法返回值的一致性
    def gradient_as_tuple(self, out_grad: "Value", node: "Value") -> Tuple["Value"]:
        """ Convenience method to always return a tuple from gradient call"""
        output = self.gradient(out_grad, node)
        if isinstance(output, tuple):
            return output
        elif isinstance(output, list):
            return tuple(output)
        else:
            return (output,)

2. TensorOp类：

# TensorOp是Op的子类，实现了父类里的__call__
class TensorOp(Op):
    """ Op class specialized to output tensors, will be alternate subclasses for other structures """

    def __call__(self, *args):
        # 如调用result=TensorOp()(tensor1, tensor 2)，就会调用TensorOp的__call__函数
        # *args也就是输入的tensors
        # 这里make_from_op是在用算子对args里的tensor做运算，返回运算结果
        # 比如Tensor.make_from_op(EWiseAdd(), *(a, b))就是把a b加在一起返回
        return Tensor.make_from_op(self, args)  #定义在Value类中

3. Value类：是Tensor的父类
   Value即计算图中的节点，如下代码可知一个节点中包含信息有：该节点的运算op、该节点的inputs、该节点的值（即inputs进行op后得到的值）、是否参与梯度计算

# Value是Tensor的父类
class Value:
    """A value in the computational graph."""

    # 用来追踪计算图
    op: Optional[Op] # 算子
    inputs: List["Value"] # 计算图中输入到这个结点的其他结点
    # 比如v3 = v1 + v2，那么v3.op就是加法算子EWiseAdd，v3.inputs是[v1, v2]
    
    # The following fields are cached fields for
    # dynamic computation
    cached_data: NDArray # 存储自身的数值
    requires_grad: bool # 是否参与梯度计算

    def realize_cached_data(self):
        """Run compute to realize the cached data"""
        # avoid recomputation
        if self.cached_data is not None:
            return self.cached_data
        # note: data implicitly calls realized cached data
        # 用自己的op算子和inputs输入算出自己的数值
        self.cached_data = self.op.compute(*[x.realize_cached_data() for x in self.inputs])
        return self.cached_data

    def is_leaf(self):
        # 叶节点是没有op的
        return self.op is None

    def __del__(self):
        global TENSOR_COUNTER
        TENSOR_COUNTER -= 1

    # 一个初始化函数，为op、inputs等赋初始值
    def _init(self, op: Optional[Op], inputs: List["Tensor"], *, num_outputs: int = 1, cached_data: List[object] = None, requires_grad: Optional[bool] = None):
        global TENSOR_COUNTER
        TENSOR_COUNTER += 1
        # any() 函数用于判断可迭代对象中是否存在至少一个为真的元素
        # 若可迭代对象中的任何一个元素为真，则返回 True；若所有元素都为假，则返回 False。
        # 如果用户没有明确指定requires_grad参数，那么根据inputs的情况来确定是否需要计算梯度
        # 如果至少有一个输入张量需要计算梯度，那么就将本节点requires_grad设为真；否则设为假。这样可以简化用户在使用时的输入，并且使得模型更加灵活。
        if requires_grad is None:
            requires_grad = any(x.requires_grad for x in inputs)
        self.op = op
        self.inputs = inputs
        self.num_outputs = num_outputs
        self.cached_data = cached_data
        self.requires_grad = requires_grad

    # 和子类Tensor里的一样
    # @classmethod让下面的函数不用实例化就可以调用，即直接Value.make_const()
    # 作用是创建了一个常量张量对象，并在创建时进行了初始化
    @classmethod
    def make_const(cls, data, *, requires_grad=False):
        value = cls.__new__(cls)  # 创建一个新的本类的实例对象（cls用法类似self，表示类本身，self表示实例本身）
        
        # 初始化这个新创建的实例，注意这里是上面定义的_init()而不是__init__()
        value._init(None, [], cached_data=data, requires_grad=requires_grad,)
        return value

    # 和子类Tensor里的一样
    @classmethod
    def make_from_op(cls, op: Op, inputs: List["Value"]):
        value = cls.__new__(cls)
        value._init(op, inputs)

        if not LAZY_MODE:
            if not value.requires_grad:
                return value.detach()
            value.realize_cached_data()
        return value

4. Tensor类：

class Tensor(Value):
    grad: "Tensor"

	# 初始化张量对象
    def __init__(self, array, *, device: Optional[Device] = None, dtype=None, requires_grad=True, **kwargs):
        if isinstance(array, Tensor):
            if device is None:
                device = array.device
            if dtype is None:
                dtype = array.dtype
            if device == array.device and dtype == array.dtype:
                cached_data = array.realize_cached_data()
            else:
                # fall back, copy through numpy conversion
                # 如果dtype或device不一致，则重新创建
                cached_data = Tensor._array_from_numpy(
                    array.numpy(), device=device, dtype=dtype
                )
        else:
            device = device if device else cpu()
            cached_data = Tensor._array_from_numpy(array, device=device, dtype=dtype)

        # 给cached_data赋值，因为不是op产生的tensor，所以op和inputs都是空的
        self._init(None, [], cached_data=cached_data, requires_grad=requires_grad,)

    # 用后端生成数组
    @staticmethod
    def _array_from_numpy(numpy_array, device, dtype):
        if array_api is numpy:
            return numpy.array(numpy_array, dtype=dtype)
        return array_api.array(numpy_array, device=device, dtype=dtype)

    # 用op和inputs计算出一个新tensor
    @staticmethod
    def make_from_op(op: Op, inputs: List["Value"]):
        tensor = Tensor.__new__(Tensor)
        tensor._init(op, inputs)
        if not LAZY_MODE:
            tensor.realize_cached_data()
        return tensor

    # 返回一个无op，无inputs，只有cached_data，且默认不加入梯度计算的value
    @staticmethod
    def make_const(data, requires_grad=False):
        tensor = Tensor.__new__(Tensor)
        tensor._init(
            None,
            [],
            cached_data=data
            if not isinstance(data, Tensor)
            else data.realize_cached_data(),
            requires_grad=requires_grad,
        )
        return tensor

    # 可以通过.data获取没有梯度的数值，梯度下降更新参数时会用到
    @property
    def data(self):
        return self.detach()

    # 更新data的方法，这里要求更新前后dtype一致
    @data.setter
    def data(self, value):
        assert isinstance(value, Tensor)
        assert value.dtype == self.dtype, "%s %s" % (
            value.dtype,
            self.dtype,
        )
        self.cached_data = value.realize_cached_data()

    # 用make_const返回一个无梯度、在计算图外的tensor
    def detach(self):
        """Create a new tensor that shares the data but detaches from the graph."""
        return Tensor.make_const(self.realize_cached_data())

    # 因为目前只有numpy后端，所以shape就是numpy里的shape
    @property
    def shape(self):
        return self.realize_cached_data().shape

    @property
    def dtype(self):
        return self.realize_cached_data().dtype

    @property
    def device(self):
        data = self.realize_cached_data()
        # numpy array always sits on cpu
        # 如果后端是numpy，则只能用cpu
        if array_api is numpy:
            return cpu()
        return data.device

    # 用backward函数反向传播梯度
    def backward(self, out_grad=None):
        out_grad = out_grad if out_grad else Tensor(numpy.ones(self.shape))
        compute_gradient_of_variables(self, out_grad)

    def __repr__(self):
        return "needle.Tensor(" + str(self.realize_cached_data()) + ")"

    def __str__(self):
        return self.realize_cached_data().__str__()

    def numpy(self):
        data = self.realize_cached_data()
        if array_api is numpy:
            return data
        return data.numpy()

    # + 重载
    def __add__(self, other):
        if isinstance(other, Tensor):
            return needle.ops.EWiseAdd()(self, other)
        else:
            return needle.ops.AddScalar(other)(self)

    # * 重载
    def __mul__(self, other):
        if isinstance(other, Tensor):
            return needle.ops.EWiseMul()(self, other)
        else:
            return needle.ops.MulScalar(other)(self)

    # ** 重载
    def __pow__(self, other):
        if isinstance(other, Tensor):
            raise NotImplementedError()
        else:
            return needle.ops.PowerScalar(other)(self)

    # - 重载
    def __sub__(self, other):
        if isinstance(other, Tensor):
            return needle.ops.EWiseAdd()(self, needle.ops.Negate()(other))
        else:
            return needle.ops.AddScalar(-other)(self)

    # / 重载
    def __truediv__(self, other):
        if isinstance(other, Tensor):
            return needle.ops.EWiseDiv()(self, other)
        else:
            return needle.ops.DivScalar(other)(self)

    # @ 重载
    def __matmul__(self, other):
        return needle.ops.MatMul()(self, other)

    # matmul也可以不用运算符重载
    def matmul(self, other):
        return needle.ops.MatMul()(self, other)

    def sum(self, axes=None):
        return needle.ops.Summation(axes)(self)

    def broadcast_to(self, shape):
        return needle.ops.BroadcastTo(shape)(self)

    def reshape(self, shape):
        return needle.ops.Reshape(shape)(self)

    # 负号重载
    def __neg__(self):
        return needle.ops.Negate()(self)

    def transpose(self, axes=None):
        return needle.ops.Transpose(axes)(self)

    __radd__ = __add__
    __rmul__ = __mul__
    __rsub__ = __sub__
    __rmatmul__ = __matmul__

各种函数

1. compute_gradient_of_variables
2. find_topo_sort
3. topo_sort_dfs