import random from typing import Iterable, Optional, Sequence, Tuple, Union
import numba from numba import cuda import numpy as np import numpy.typing as npt from numpy import array, float64 from typing_extensions import TypeAlias
from .operators import prod
MAX_DIMS = 32
classIndexingError(RuntimeError): "Exception raised for indexing errors." pass
defpermute(self, *order: int) -> TensorData: """ Permute the dimensions of the tensor.
Args: *order: a permutation of the dimensions
Returns: New `TensorData` with the same storage and a new dimension order. """ assertlist(sorted(order)) == list( range(len(self.shape)) ), f"Must give a position to each dimension. Shape: {self.shape} Order: {order}"
# 以 permute 方法为例，说明 shape & strides 表示法优点： # 进行 permute 操作时无任何 storage 层拷贝，只需要根据传入参数 order 重新排列 shape & strides 即可 return TensorData( self._storage, tuple([self.shape[o] for o in order]), tuple([self._strides[o] for o in order]), )
# 该函数核心功能：将 element-wise fn 封装为一个接受 Tensor 的 Storage/Shape/Strides 作为参数的可调用对象 # 值得一提的是，这里的可调用对象参数要求是 np.array，而不是 Python 原生类型 deftensor_map( fn: Callable[[float], float] ) -> Callable[[Storage, Shape, Strides, Storage, Shape, Strides], None]: """ Low-level implementation of tensor map between tensors with *possibly different strides*.
Simple version:
* Fill in the `out` array by applying `fn` to each value of `in_storage` assuming `out_shape` and `in_shape` are the same size.
Broadcasted version:
* Fill in the `out` array by applying `fn` to each value of `in_storage` assuming `out_shape` and `in_shape` broadcast. (`in_shape` must be smaller than `out_shape`).
Args: fn: function from float-to-float to apply
Returns: Tensor map function. """
def_map( out: Storage, out_shape: Shape, out_strides: Strides, in_storage: Storage, in_shape: Shape, in_strides: Strides, ) -> None: out_index: Index = np.zeros(MAX_DIMS, np.int16) in_index: Index = np.zeros(MAX_DIMS, np.int16) for i inrange(len(out)): # 处理 index 映射广播等 to_index(i, out_shape, out_index) broadcast_index(out_index, out_shape, in_shape, in_index) o = index_to_position(out_index, out_strides) j = index_to_position(in_index, in_strides) # 调用传入的 fn 函数 out[o] = fn(in_storage[j])
classTensorBackend: def__init__(self, ops: Type[TensorOps]): """ Dynamically construct a tensor backend based on a `tensor_ops` object that implements map, zip, and reduce higher-order functions.
Args: ops : tensor operations object see `tensor_ops.py`
defaccumulate_derivative(self, x: Any) -> None: """ Add `val` to the the derivative accumulated on this variable. Should only be called during autodifferentiation on leaf variables.
Args: x : value to be accumulated """ # 叶子节点梯度累加 assert self.is_leaf(), "Only leaf variables can have derivatives." if self.grad isNone: self.grad = Tensor.make( [0] * int(operators.prod(self.shape)), self.shape, backend=self.backend ) self.grad += x
defis_leaf(self) -> bool: "True if this variable created by the user (no `last_fn`)" return self.history isnotNoneand self.history.last_fn isNone
defsave_for_backward(self, *values: Any) -> None: "Store the given `values` if they need to be used during backpropagation." if self.no_grad: return self.saved_values = values
@classmethod defapply(cls, *vals: Tensor) -> Tensor: raw_vals = [] need_grad = False for v in vals: if v.requires_grad(): need_grad = True raw_vals.append(v.detach())
# Create the context. ctx = Context(not need_grad)
# Call forward with the variables. c = cls._forward(ctx, *raw_vals) # assert isinstance(c, Tensor), "Expected return type Tensor got %s" % ( # type(c) # )
# Create a new variable from the result with a new history. back = None if need_grad: back = minitorch.History(cls, ctx, vals) return minitorch.Tensor(c._tensor, back, backend=c.backend)
classMul(Function): @staticmethod defforward(ctx: Context, a: Tensor, b: Tensor) -> Tensor: ctx.save_for_backward(a, b) return a.f.mul_zip(a, b)
@staticmethod defbackward(ctx: Context, grad_output: Tensor) -> Tuple[Tensor, Tensor]: a, b = ctx.saved_values return ( grad_output.f.mul_zip(b, grad_output), grad_output.f.mul_zip(a, grad_output), )
classVariable(Protocol): defaccumulate_derivative(self, x: Any) -> None: """ Accumulates the derivative (gradient) for this Variable.
Args: x (Any): The gradient value to be accumulated. """ pass
@property defunique_id(self) -> int: """ Returns: int: The unique identifier of this Variable. """ pass
defis_leaf(self) -> bool: """ Returns whether this Variable is a leaf node in the computation graph.
Returns: bool: True if this Variable is a leaf node, False otherwise. """ pass
defis_constant(self) -> bool: """ Returns whether this Variable represents a constant value.
Returns: bool: True if this Variable is constant, False otherwise. """ pass
@property defparents(self) -> Iterable["Variable"]: """ Returns the parent Variables of this Variable in the computation graph.
Returns: Iterable[Variable]: The parent Variables of this Variable. """ pass
defchain_rule(self, d_output: Any) -> Iterable[Tuple["Variable", Any]]: """ Implements the chain rule to compute the gradient contributions of this Variable.
Args: d_output (Any): The gradient of the output with respect to the Variable.
Returns: Iterable[Tuple[Variable, Any]]: An iterable of tuples, where each tuple contains a parent Variable and the corresponding gradient contribution. """ pass
deftopological_sort(variable: Variable) -> Iterable[Variable]: """ Computes the topological order of the computation graph.
Args: variable: The right-most variable
Returns: Non-constant Variables in topological order starting from the right.
Hints: 1. Ensure that you visit the computation graph in a post-order depth-first search. 2. When the children nodes of the current node are visited, add the current node at the front of the result order list. """ # BEGIN HW2_1 # TODO
# END HW2_1
defbackpropagate(variable: Variable, deriv: Any) -> None: """ Runs backpropagation on the computation graph in order to compute derivatives for the leaf nodes.
Args: variable: The right-most variable deriv : Its derivative that we want to propagate backward to the leaves.
No return. Should write its results to the derivative values of each leaf through `accumulate_derivative`.
Hints: 1. Traverse nodes in topological order 2. If the node is a leaf, the derivative should be accumulated 3. Otherwise, the derivative should be propagated via chain rule """ # BEGIN HW2_1 # TODO
# END HW2_1
可以看到：
class Variable(Protocol) 定义了
autodiff.py 模块需要用到的几个方法。可以结合
class Tensor 的实现理解，这里简单说明下：
is_constant()：表示节点是否是常量。
is_leaf()：表示节点是否是
Parameter。Tensor
内部实现判断逻辑为：self.history is not None and self.history.last_fn is None。构造
Parameter 对象时，内部会调用 Tensor 的
requires_grad_
方法：self.history = History() if require_grad else None，让一个
Tensor
变成叶子节点。注意，这里的叶子节点概念不同于数据结构中树/图的叶子节点（DAG
表示中的叶子节点），而是既要是 DAG
中的叶子节点，又要是可学习参数才可以（训练输入数据和标签 x/y 在 DAG
角度也是叶子节点，但是 Tensor 角度而言其是 is_constant()
而不是 is_leaf()）。
accumulate_derivative：用于叶子节点梯度累加。
parents：获取当前 Tensor 的输入 Tensor。
chain_rule：
topological_sort 接受 forward 形成的 DAG 的最后一个
Tensor 作为参数，以深度优先搜索方式遍历 DAG（同时需要过滤掉
is_constant() 的不需要求梯度的节点），形成一个拓扑序的
Tensor 列表。需要注意的是，最后返回 Tensor 列表时需要
reverse 下，因为 backpropagate
内部进行反向传播梯度计算时，要从 DAG 的最后一个节点（即绝大多少情况下的
loss 节点）反向进行。
backpropagate 接受 DAG 最后一个节点（loss
节点）和上游传递下来的梯度（从只有一个元素的 loss Tensor
调用 backward() 时，传入
Tensor.make([1.0], (1,), backend=self.backend) ，即梯度
1.0）。其内部实现流程大概如下：
调用 topological_sort 函数获取从后向前的 Tensor
列表
定义一个 dict，保存输入 Tensor id 到输入 grad 的映射关系
遍历列表中的每一个节点
如果该节点为叶子节点，进行梯度累加
否则（即对于计算图的中间节点），调用节点的 chain_rule
方法进行链式梯度向父节点传播，保存 chain_rule 返回的非
is_constant() 节点的 Tensor id -> grad 映射到 dict
供后续使用。
值得一提的是，backpropagate
的实现对叶子节点（Parameter）与 DAG
的中间节点采用了两套截然不同的处理逻辑：
对叶子节点，进行梯度累加。
对中间节点，使用适当数据结构临时保存其 grad。
backpropagate 调用完成后，中间节点的梯度值不复存在（被
gc 回收掉）；而叶子节点的梯度被累加到 class Tensor 的
self.grad 成员上。
# 反向传播时当 input tensor 与 grad tensor shape 不一致时处理广播的辅助函数 defexpand(self, other: Tensor) -> Tensor: """ Method used to allow for backprop over broadcasting. This method is called when the output of `backward` is a different size than the input of `forward`.
Parameters: other : backward tensor (must broadcast with self)
Returns: Expanded version of `other` with the right derivatives
"""
# Case 1: Both the same shape. if self.shape == other.shape: return other
# Case 2: Backward is a smaller than self. Broadcast up. true_shape = TensorData.shape_broadcast(self.shape, other.shape) buf = self.zeros(true_shape) self.backend.id_map(other, buf) if self.shape == true_shape: return buf
# Case 3: Still different, reduce extra dims. out = buf orig_shape = [1] * (len(out.shape) - len(self.shape)) + list(self.shape) for dim, shape inenumerate(out.shape): if orig_shape[dim] == 1and shape != 1: out = self.backend.add_reduce(out, dim) assert out.size == self.size, f"{out.shape}{self.shape}" # START CODE CHANGE (2021) return Tensor.make(out._tensor._storage, self.shape, backend=self.backend) # END CODE CHANGE (2021)
classModule: """ Modules form a tree that store parameters and other submodules. They make up the basis of neural network stacks.
Attributes: _modules : Storage of the child modules _parameters : Storage of the module's parameters training : Whether the module is in training mode or evaluation mode
defmodules(self) -> Sequence[Module]: "Return the direct child modules of this module." m: Dict[str, Module] = self.__dict__["_modules"] returnlist(m.values())
# 设置训练/推理模式，需要对子模块递归设置 deftrain(self) -> None: "Set the mode of this module and all descendent modules to `train`." # ASSIGN0.4 for m in self.modules(): m.train() self.training = True # END ASSIGN0.4
defeval(self) -> None: "Set the mode of this module and all descendent modules to `eval`." for m in self.modules(): m.eval() self.training = False
defnamed_parameters(self) -> Sequence[Tuple[str, Parameter]]: """ Collect all the parameters of this module and its descendents.
Returns: The name and `Parameter` of each ancestor parameter. """
# Collect our parameters and give them a name. parameters = {} for k, v in self._parameters.items(): parameters[k] = v
# 递归访问子 Module 获取子模块参数 # Recurse down to children submodules for mod_name, m in self._modules.items(): for k, v in m.named_parameters(): parameters[f"{mod_name}.{k}"] = v returnlist(parameters.items())
defparameters(self) -> Sequence[Parameter]: "Enumerate over all the parameters of this module and its descendents." return [j for _, j in self.named_parameters()]
defadd_parameter(self, k: str, v: Any) -> Parameter: """ Manually add a parameter. Useful helper for scalar parameters.
Args: k: Local name of the parameter. v: Value for the parameter.
Returns: Newly created parameter. """ val = Parameter(v, k) self.__dict__["_parameters"][k] = val return val
defzero_grad(self) -> None: for p in self.parameters: if p.value isNone: continue ifhasattr(p.value, "derivative"): if p.value.derivative isnotNone: p.value.derivative = None ifhasattr(p.value, "grad"): if p.value.grad isnotNone: p.value.grad = None
defstep(self) -> None: for p in self.parameters: if p.value isNone: continue ifhasattr(p.value, "grad"): if p.value.grad isnotNone: # 梯度下降更新 Parameter p.update(p.value - self.lr * p.value.grad)
def_print(self) -> None: for param in self.parameters: if param.value isNone: continue print(param.value.shape) print(param.value.grad)
