"""
tensor.py —— 一个最小的张量深度学习框架（NumPy 版 miniTorch）：带步长的存储、可替换的后端、张量自动微分、nn 模块、优化器
=====================================================================
（第 5 讲；对应原课 HW1 的"后端"、HW2 的自动微分与网络。第 6 讲之后的 Transformer 也可以用它搭）

【这个文件要回答的问题】
    一个深度学习框架由哪几层组成？每一层只管自己的事：
        ┌──────────────────────────────────────────────────────────────┐
        │ nn / optim  : Module、Parameter、Linear、Embedding、SGD、Adam  │ ← 写模型、训练的人用
        ├──────────────────────────────────────────────────────────────┤
        │ Tensor      : 运算符重载、计算图（History）、backward()        │ ← 自动微分（复用 autodiff.py）
        ├──────────────────────────────────────────────────────────────┤
        │ Function    : 每个运算的 forward / backward（局部导数）        │
        ├──────────────────────────────────────────────────────────────┤
        │ 后端 Ops    : map、zip、reduce、matmul 四类底层运算             │ ← 第 2–4 讲的内核在这一层
        │   PythonOps（纯 Python 循环）/ NumpyOps / SimtOps（第 3 讲的模拟器内核）│
        ├──────────────────────────────────────────────────────────────┤
        │ TensorData  : 一维存储 + 形状 + 步长                           │ ← P1 §5
        └──────────────────────────────────────────────────────────────┘
    换一个后端，上面几层一行都不用改——这就是"正确的抽象"（第 1 讲 §4.3）。

【本文件提供】
    TensorData                       : 存储 + 形状 + 步长；permute、to_numpy、is_contiguous
    NumpyOps、PythonOps、SimtOps     : 三个后端，接口相同：map(名字)、zip(名字)、reduce(名字)、matmul
    Tensor、tensor()、zeros()、ones() : 张量与创建函数；运算：+ − × ÷ @、exp、log、relu、sigmoid、sum、mean、
                                       permute、view、contiguous、embedding
    Function 及其子类                 : Neg、Add、Mul、Inv、Exp、Log、ReLU、Sigmoid、Sum、MatMul、Permute、View、Copy、Embedding
    Parameter、Module、Linear、Embedding、SGD、Adam

【约定】
    - 数据一律 float64（梯度检查要用双精度，slides 第 34 页）；SimtOps 内部用 float32（第 3 讲的内核是 float32 的）。
    - 和 miniTorch 一样：标量用形状 (1,) 的张量表示；sum(dim) 保留被归约的那一维（长度变成 1）。
    - 广播：Add、Mul 等运算的输入形状可以不同；反向时梯度会自动"求和缩回"输入的形状（Tensor.chain_rule 里的 _sum_to）。
"""

import math
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "prereq1"))
from autodiff import Context, backpropagate  # noqa: E402
from tensor_layout import broadcast_index, index_to_position, shape_broadcast, strides_for, to_index  # noqa: E402

_next_id = [10 ** 9]          # 和 autodiff.Scalar 的编号错开


def _new_id():
    _next_id[0] += 1
    return _next_id[0]


# =====================================================================
# 1. 存储：TensorData
# =====================================================================

class TensorData:
    """
    一维存储 storage（float64）+ 形状 shape + 步长 strides（单位是元素）。和第 3 讲的 HostTensor、miniTorch 的 TensorData 相同。

    【方法】permute(*order)：交换维度（零拷贝）；to_numpy()：按步长读出一个普通数组；is_contiguous()
    """

    def __init__(self, storage, shape, strides=None):
        self.storage = storage
        self.shape = tuple(int(s) for s in shape)
        self.strides = tuple(strides) if strides is not None else strides_for(self.shape)
        self.size = int(np.prod(self.shape)) if self.shape else 1

    @classmethod
    def from_numpy(cls, arr):
        arr = np.ascontiguousarray(arr, dtype=np.float64)
        return cls(arr.reshape(-1).copy(), arr.shape)

    def view_numpy(self):
        """不复制，按 shape / strides 得到一个 NumPy 视图（只读用）。"""
        item = self.storage.itemsize
        return np.lib.stride_tricks.as_strided(self.storage, self.shape, tuple(s * item for s in self.strides),
                                               writeable=False)

    def to_numpy(self):
        return self.view_numpy().copy()

    def permute(self, *order):
        return TensorData(self.storage, tuple(self.shape[k] for k in order), tuple(self.strides[k] for k in order))

    def is_contiguous(self):
        want = strides_for(self.shape)
        return all(n == 1 or s == w for n, s, w in zip(self.shape, self.strides, want))


# =====================================================================
# 2. 后端：四类底层运算
# =====================================================================

def _sig_np(x):
    """数值稳定的 sigmoid（x 很负时不让 exp(−x) 溢出）。"""
    out = np.empty_like(x)
    pos = x >= 0
    out[pos] = 1.0 / (1.0 + np.exp(-x[pos]))
    e = np.exp(x[~pos])
    out[~pos] = e / (1.0 + e)
    return out


def _sig_py(x):
    return 1.0 / (1.0 + math.exp(-x)) if x >= 0 else math.exp(x) / (1.0 + math.exp(x))


# 名字 → (NumPy 向量化版本, 标量版本)
MAP_FNS = {
    "id": (lambda x: x.copy(), lambda x: x),
    "neg": (np.negative, lambda x: -x),
    "inv": (np.reciprocal, lambda x: 1.0 / x),
    "exp": (np.exp, math.exp),
    "log": (np.log, math.log),
    "relu": (lambda x: np.maximum(x, 0.0), lambda x: x if x > 0 else 0.0),
    "sigmoid": (_sig_np, _sig_py),
}
ZIP_FNS = {
    "add": (np.add, lambda a, b: a + b),
    "mul": (np.multiply, lambda a, b: a * b),
    "relu_back": (lambda x, g: np.where(x > 0, g, 0.0), lambda x, g: g if x > 0 else 0.0),   # ReLU 的反向：x>0 处放行梯度
}
REDUCE_FNS = {
    "sum": (np.add, lambda a, b: a + b, 0.0),
    "max": (np.maximum, lambda a, b: a if a > b else b, -math.inf),
}


class NumpyOps:
    """后端一：用 NumPy 的向量化运算（对应真实框架里调用优化过的库）。"""
    name = "numpy"

    @staticmethod
    def map(fn):
        f = MAP_FNS[fn][0]
        return lambda a: TensorData.from_numpy(f(a.view_numpy()))

    @staticmethod
    def zip(fn):
        f = ZIP_FNS[fn][0]
        return lambda a, b: TensorData.from_numpy(f(a.view_numpy(), b.view_numpy()))

    @staticmethod
    def reduce(fn):
        f = REDUCE_FNS[fn][0]
        return lambda a, dim: TensorData.from_numpy(f.reduce(a.view_numpy(), axis=dim, keepdims=True))

    @staticmethod
    def matmul(a, b):
        return TensorData.from_numpy(np.matmul(a.view_numpy(), b.view_numpy()))


class PythonOps:
    """
    后端二：纯 Python 循环，每个元素都用 to_index / broadcast_index / index_to_position 找位置。
    它就是第 3 讲那些内核"去掉线程编号"之后的样子（原课 miniTorch 的 CPU 参考实现也是这样写的）。很慢，只用于小张量。
    """
    name = "python"

    @staticmethod
    def map(fn):
        f = MAP_FNS[fn][1]

        def run(a):
            out = np.zeros(a.size)
            st = strides_for(a.shape)
            for i in range(a.size):
                idx = to_index(i, a.shape)
                out[index_to_position(idx, st)] = f(a.storage[index_to_position(idx, a.strides)])
            return TensorData(out, a.shape)
        return run

    @staticmethod
    def zip(fn):
        f = ZIP_FNS[fn][1]

        def run(a, b):
            shape = shape_broadcast(a.shape, b.shape)
            size = int(np.prod(shape))
            out = np.zeros(size)
            for i in range(size):
                idx = to_index(i, shape)
                ai = broadcast_index(idx, shape, a.shape)
                bi = broadcast_index(idx, shape, b.shape)
                out[i] = f(a.storage[index_to_position(ai, a.strides)], b.storage[index_to_position(bi, b.strides)])
            return TensorData(out, shape)
        return run

    @staticmethod
    def reduce(fn):
        _, f, start = REDUCE_FNS[fn]

        def run(a, dim):
            shape = tuple(1 if k == dim else s for k, s in enumerate(a.shape))
            size = int(np.prod(shape))
            out = np.zeros(size)
            for i in range(size):
                idx = list(to_index(i, shape))
                acc = start
                for k in range(a.shape[dim]):
                    idx[dim] = k
                    acc = f(acc, a.storage[index_to_position(idx, a.strides)])
                out[i] = acc
            return TensorData(out, shape)
        return run

    @staticmethod
    def matmul(a, b):
        A, B = a.to_numpy(), b.to_numpy()
        squeeze = A.ndim == 2 and B.ndim == 2
        A3 = A if A.ndim == 3 else A[None]
        B3 = B if B.ndim == 3 else B[None]
        nb = max(A3.shape[0], B3.shape[0])
        n, m, p = A3.shape[1], A3.shape[2], B3.shape[2]
        out = np.zeros((nb, n, p))
        for z in range(nb):
            Az = A3[z if A3.shape[0] > 1 else 0]
            Bz = B3[z if B3.shape[0] > 1 else 0]
            for i in range(n):
                for j in range(p):
                    acc = 0.0
                    for k in range(m):
                        acc += Az[i, k] * Bz[k, j]
                    out[z, i, j] = acc
        return TensorData.from_numpy(out[0] if squeeze else out)


class SimtOps:
    """
    后端三：调用第 3 讲在 SIMT 模拟器上写的内核（kernels.py）。数据在 float64 和内核用的 float32 之间转换。
    只用来演示"后端可以换"——它比纯 Python 还慢，结果的精度是 float32 的。
    """
    name = "simt"
    _dev = None

    @classmethod
    def _k(cls):
        sys.path.insert(0, os.path.join(HERE, "..", "lecture03"))
        import kernels  # noqa: E402
        from simt import Device  # noqa: E402
        if cls._dev is None:
            cls._dev = Device()
        return kernels, cls._dev

    @staticmethod
    def _host(a, kernels):
        return kernels.HostTensor(a.storage.astype(np.float32), a.shape, a.strides)

    @classmethod
    def map(cls, fn):
        f = MAP_FNS[fn][1]

        def run(a):
            k, dev = cls._k()
            out, _ = k.tensor_map(dev, f, cls._host(a, k))
            return TensorData.from_numpy(out.to_numpy())
        return run

    @classmethod
    def zip(cls, fn):
        f = ZIP_FNS[fn][1]

        def run(a, b):
            k, dev = cls._k()
            out, _ = k.tensor_zip(dev, f, cls._host(a, k), cls._host(b, k))
            return TensorData.from_numpy(out.to_numpy())
        return run

    @classmethod
    def reduce(cls, fn):
        _, f, start = REDUCE_FNS[fn]

        def run(a, dim):
            k, dev = cls._k()
            out, _ = k.tensor_reduce(dev, f, cls._host(a, k), dim, start, method="block")
            return TensorData.from_numpy(out.to_numpy())
        return run

    @classmethod
    def matmul(cls, a, b):
        k, dev = cls._k()
        out, _ = k.tensor_matmul(dev, cls._host(a, k), cls._host(b, k))
        return TensorData.from_numpy(out.to_numpy())


# =====================================================================
# 3. 张量与自动微分
# =====================================================================

class History:
    """张量的出身：last_fn（哪个 Function）、ctx（前向保存的值）、inputs（输入张量）。叶子的 last_fn 为 None。"""

    def __init__(self, last_fn=None, ctx=None, inputs=()):
        self.last_fn = last_fn
        self.ctx = ctx
        self.inputs = inputs


class Tensor:
    """
    参与自动微分的张量。实现了 autodiff.py 的"变量协议"，所以 backpropagate 可以直接处理它。

    【创建】tensor(数组, requires_grad=True) / zeros(形状) / ones(形状)
    【属性】shape、size、grad（反向之后的梯度，也是 Tensor）、backend
    【边界】只有 history 不为 None 的张量才会被记录进计算图；requires_grad=False 的张量是常量。
    """

    def __init__(self, data, backend=NumpyOps, history=None, name=None):
        self.data = data
        self.backend = backend
        self.history = history
        self.grad = None
        self.unique_id = _new_id()
        self.name = name

    # ---- 基本信息 ----
    @property
    def shape(self):
        return self.data.shape

    @property
    def size(self):
        return self.data.size

    def to_numpy(self):
        return self.data.to_numpy()

    def item(self):
        assert self.size == 1, "item() 只能用于只有一个元素的张量"
        return float(self.data.view_numpy().reshape(-1)[0])

    def __repr__(self):
        return f"Tensor(shape={self.shape}, backend={self.backend.name})"

    # ---- 变量协议 ----
    def requires_grad_(self, flag=True):
        self.history = History() if flag else None
        return self

    def is_leaf(self):
        return self.history is not None and self.history.last_fn is None

    def is_constant(self):
        return self.history is None

    @property
    def parents(self):
        return self.history.inputs if self.history is not None else ()

    def chain_rule(self, d_output):
        """调用 last_fn.backward 得到每个输入的梯度；如果输入被广播过，把梯度求和缩回输入的形状。"""
        h = self.history
        grads = h.last_fn.backward(h.ctx, d_output)
        return [(inp, _sum_to(g, inp.shape)) for inp, g in zip(h.inputs, grads)]

    def accumulate_derivative(self, d):
        self.grad = d if self.grad is None else self.grad + d

    def backward(self, grad_output=None):
        """反向传播。grad_output 默认是和自己同形状的全 1（损失是标量时就是 d loss / d loss = 1）。"""
        if grad_output is None:
            grad_output = Tensor(TensorData.from_numpy(np.ones(self.shape)), self.backend)
        backpropagate(self, grad_output)

    def zero_grad_(self):
        self.grad = None

    def detach(self):
        """同一份数据、但不带历史的常量张量。"""
        return Tensor(self.data, self.backend)

    # ---- 运算 ----
    def _ensure(self, b):
        return b if isinstance(b, Tensor) else Tensor(TensorData.from_numpy(np.array([float(b)])), self.backend)

    def __add__(self, b):
        return Add.apply(self, self._ensure(b))

    __radd__ = __add__

    def __mul__(self, b):
        return Mul.apply(self, self._ensure(b))

    __rmul__ = __mul__

    def __neg__(self):
        return Neg.apply(self)

    def __sub__(self, b):
        return Add.apply(self, Neg.apply(self._ensure(b)))

    def __rsub__(self, b):
        return Add.apply(self._ensure(b), Neg.apply(self))

    def __truediv__(self, b):
        return Mul.apply(self, Inv.apply(self._ensure(b)))

    def __rtruediv__(self, b):
        return Mul.apply(self._ensure(b), Inv.apply(self))

    def __matmul__(self, b):
        return MatMul.apply(self, b)

    def exp(self):
        return Exp.apply(self)

    def log(self):
        return Log.apply(self)

    def relu(self):
        return ReLU.apply(self)

    def sigmoid(self):
        return Sigmoid.apply(self)

    def sum(self, dim=None):
        """沿 dim 求和（保留这一维）；dim=None 时把所有元素加成形状 (1,) 的张量。"""
        if dim is None:
            return Sum.apply(self.contiguous().view(self.size), 0)
        return Sum.apply(self, dim)

    def mean(self, dim=None):
        n = self.size if dim is None else self.shape[dim]
        return self.sum(dim) * (1.0 / n)

    def permute(self, *order):
        return Permute.apply(self, order)

    def transpose_last2(self):
        order = list(range(len(self.shape)))
        order[-1], order[-2] = order[-2], order[-1]
        return self.permute(*order)

    def contiguous(self):
        return self if self.data.is_contiguous() else Copy.apply(self)

    def view(self, *shape):
        return View.apply(self, shape)


def tensor(values, requires_grad=False, backend=NumpyOps):
    """从 NumPy 数组 / 列表创建张量。"""
    t = Tensor(TensorData.from_numpy(np.asarray(values, dtype=np.float64)), backend)
    return t.requires_grad_() if requires_grad else t


def zeros(shape, backend=NumpyOps):
    return Tensor(TensorData.from_numpy(np.zeros(shape)), backend)


def ones(shape, backend=NumpyOps):
    return Tensor(TensorData.from_numpy(np.ones(shape)), backend)


def _sum_to(g, shape):
    """
    广播的反向：前向时输入被"复制"到了更大的形状，反向时梯度要把复制出来的那些份加回去。
        多出来的前导维度 → 沿它求和后去掉；输入里长度为 1、输出里大于 1 的维度 → 沿它求和（保留长度 1）。
    【例子】偏置 b 的形状 (h,)，加到 (B, h) 上；b 的梯度 = 输出梯度沿第 0 维求和。
    """
    if g.shape == tuple(shape):
        return g
    ndiff = len(g.shape) - len(shape)
    for k in range(ndiff):
        g = Sum.apply(g, 0).view(*g.shape[1:])
    for k, s in enumerate(shape):
        if s == 1 and g.shape[k] != 1:
            g = Sum.apply(g, k)
    return g


class Function:
    """
    所有张量运算的基类。子类写两个静态方法：
        forward(ctx, *inputs) → 输出张量（输入是不带历史的常量，所以这里的计算不会进入计算图）
        backward(ctx, grad)   → 每个张量输入的梯度（元组，顺序和张量输入一致；同样用常量张量计算）
    apply 负责：把输入去掉历史后调用 forward，然后在输出上记下 History（如果有输入需要梯度）。
    非张量参数（例如 sum 的 dim、permute 的 order、embedding 的下标）照常传给 forward，但不进入 inputs。
    """

    @classmethod
    def apply(cls, *args):
        tensors = tuple(a for a in args if isinstance(a, Tensor))
        ctx = Context()
        raw = [a.detach() if isinstance(a, Tensor) else a for a in args]
        out = cls.forward(ctx, *raw)
        if any(not t.is_constant() for t in tensors):
            out.history = History(cls, ctx, tensors)
        return out


def _new(data, like):
    return Tensor(data, like.backend)


class Neg(Function):
    @staticmethod
    def forward(ctx, a):
        return _new(a.backend.map("neg")(a.data), a)

    @staticmethod
    def backward(ctx, g):
        return (-g,)


class Inv(Function):
    @staticmethod
    def forward(ctx, a):
        ctx.save_for_backward(a)
        return _new(a.backend.map("inv")(a.data), a)

    @staticmethod
    def backward(ctx, g):
        (a,) = ctx.saved_values
        return (-(g * (1.0 / (a * a))),)


class Add(Function):
    @staticmethod
    def forward(ctx, a, b):
        return _new(a.backend.zip("add")(a.data, b.data), a)

    @staticmethod
    def backward(ctx, g):
        return g, g


class Mul(Function):
    @staticmethod
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return _new(a.backend.zip("mul")(a.data, b.data), a)

    @staticmethod
    def backward(ctx, g):
        a, b = ctx.saved_values
        return g * b, g * a


class Exp(Function):
    @staticmethod
    def forward(ctx, a):
        out = _new(a.backend.map("exp")(a.data), a)
        ctx.save_for_backward(out)
        return out

    @staticmethod
    def backward(ctx, g):
        (out,) = ctx.saved_values
        return (g * out,)


class Log(Function):
    @staticmethod
    def forward(ctx, a):
        ctx.save_for_backward(a)
        return _new(a.backend.map("log")(a.data), a)

    @staticmethod
    def backward(ctx, g):
        (a,) = ctx.saved_values
        return (g / a,)


class ReLU(Function):
    @staticmethod
    def forward(ctx, a):
        ctx.save_for_backward(a)
        return _new(a.backend.map("relu")(a.data), a)

    @staticmethod
    def backward(ctx, g):
        (a,) = ctx.saved_values
        return (_new(a.backend.zip("relu_back")(a.data, g.data), a),)


class Sigmoid(Function):
    @staticmethod
    def forward(ctx, a):
        out = _new(a.backend.map("sigmoid")(a.data), a)
        ctx.save_for_backward(out)
        return out

    @staticmethod
    def backward(ctx, g):
        (s,) = ctx.saved_values
        return (g * s * (1.0 - s),)


class Sum(Function):
    @staticmethod
    def forward(ctx, a, dim):
        ctx.save_for_backward(a.shape)
        return _new(a.backend.reduce("sum")(a.data, dim), a)

    @staticmethod
    def backward(ctx, g):
        (shape,) = ctx.saved_values
        return (g + zeros(shape, g.backend),)          # 把长度为 1 的那一维广播回原来的长度


class MatMul(Function):
    @staticmethod
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return _new(a.backend.matmul(a.data, b.data), a)

    @staticmethod
    def backward(ctx, g):
        a, b = ctx.saved_values
        return g @ b.transpose_last2(), a.transpose_last2() @ g     # dA = dC·Bᵀ，dB = Aᵀ·dC


class Permute(Function):
    @staticmethod
    def forward(ctx, a, order):
        ctx.save_for_backward(order)
        return _new(a.data.permute(*order), a)

    @staticmethod
    def backward(ctx, g):
        (order,) = ctx.saved_values
        inv = [0] * len(order)
        for i, k in enumerate(order):
            inv[k] = i
        return (g.permute(*inv),)


class View(Function):
    @staticmethod
    def forward(ctx, a, shape):
        if not a.data.is_contiguous():
            raise ValueError("view 只能用于连续存放的张量；先调用 .contiguous()")
        shape = tuple(int(s) for s in shape)
        if int(np.prod(shape)) != a.size:
            raise ValueError(f"view 的元素个数不匹配：{a.shape} → {shape}")
        ctx.save_for_backward(a.shape)
        return _new(TensorData(a.data.storage, shape), a)

    @staticmethod
    def backward(ctx, g):
        (shape,) = ctx.saved_values
        return (g.contiguous().view(*shape),)


class Copy(Function):
    """按行优先重新排一份（contiguous）。"""
    @staticmethod
    def forward(ctx, a):
        return _new(a.backend.map("id")(a.data), a)

    @staticmethod
    def backward(ctx, g):
        return (g,)


class Embedding(Function):
    """
    查表：out[b, t, :] = W[ids[b, t], :]。前向是 gather，反向是 scatter-add（同一个词出现几次，梯度就加几次）。
    这是纯数据搬运的算子，这里不经过后端，直接用 NumPy（ids 是 NumPy 整数数组，不参与求导）。
    """
    @staticmethod
    def forward(ctx, W, ids):
        ctx.save_for_backward(ids, W.shape)
        return _new(TensorData.from_numpy(W.to_numpy()[ids]), W)

    @staticmethod
    def backward(ctx, g):
        ids, shape = ctx.saved_values
        G = np.zeros(shape)
        np.add.at(G, ids, g.to_numpy())
        return (_new(TensorData.from_numpy(G), g),)


# =====================================================================
# 4. nn 与优化器
# =====================================================================

class Parameter:
    """一个可训练的参数：包着一个 requires_grad 的叶子张量 value。"""

    def __init__(self, array, name=None, backend=NumpyOps):
        self.value = tensor(array, requires_grad=True, backend=backend)
        self.name = name

    def update(self, new_array):
        """优化器用：原地替换数值（不经过计算图）。"""
        self.value.data = TensorData.from_numpy(new_array)
        self.value.grad = None


class Module:
    """
    模块的基类：parameters() 递归收集自己和子模块的 Parameter；train() / eval() 切换模式（影响 Dropout）。
    【Python 小知识】vars(self) 返回对象的属性字典，这里用它找出所有是 Parameter 或 Module 的属性。
    """

    training = True

    def parameters(self):
        out = []
        for v in vars(self).values():
            if isinstance(v, Parameter):
                out.append(v)
            elif isinstance(v, Module):
                out.extend(v.parameters())
        return out

    def train(self):
        self._set_mode(True)

    def eval(self):
        self._set_mode(False)

    def _set_mode(self, flag):
        self.training = flag
        for v in vars(self).values():
            if isinstance(v, Module):
                v._set_mode(flag)

    def __call__(self, *args):
        return self.forward(*args)


class Linear(Module):
    """
    全连接层 y = x · W + b，W 的形状 (in, out)，b 的形状 (out,)（原课 HW2：二维矩阵权重 + 一维偏置）。
    初始化：均匀分布 U(−1/√in, 1/√in)（和 PyTorch nn.Linear 的默认值相同的范围）。
    """

    def __init__(self, in_size, out_size, rng):
        bound = 1.0 / math.sqrt(in_size)
        self.W = Parameter(rng.uniform(-bound, bound, size=(in_size, out_size)), "W")
        self.b = Parameter(rng.uniform(-bound, bound, size=(out_size,)), "b")

    def forward(self, x):
        return x @ self.W.value + self.b.value


class SGD:
    """随机梯度下降：p ← p − lr · grad。"""

    def __init__(self, params, lr):
        self.params, self.lr = params, lr

    def zero_grad(self):
        for p in self.params:
            p.value.grad = None

    def step(self):
        for p in self.params:
            if p.value.grad is not None:
                p.update(p.value.to_numpy() - self.lr * p.value.grad.to_numpy())


class Adam:
    """Adam（P3 §3）：一阶、二阶矩的指数平均 + 偏差修正。公式与 torch.optim.Adam 相同。"""

    def __init__(self, params, lr, betas=(0.9, 0.999), eps=1e-8):
        self.params, self.lr, self.betas, self.eps = params, lr, betas, eps
        self.m = [np.zeros(p.value.shape) for p in params]
        self.v = [np.zeros(p.value.shape) for p in params]
        self.t = 0

    def zero_grad(self):
        for p in self.params:
            p.value.grad = None

    def step(self):
        self.t += 1
        b1, b2 = self.betas
        bc1, bc2 = 1 - b1 ** self.t, 1 - b2 ** self.t
        for i, p in enumerate(self.params):
            if p.value.grad is None:
                continue
            g = p.value.grad.to_numpy()
            self.m[i] = b1 * self.m[i] + (1 - b1) * g
            self.v[i] = b2 * self.v[i] + (1 - b2) * g * g
            denom = np.sqrt(self.v[i]) / math.sqrt(bc2) + self.eps
            p.update(p.value.to_numpy() - (self.lr / bc1) * self.m[i] / denom)