qat-int2-cifar

    def backward(self, grad: Optional[np.ndarray] = None):
        """Reverse-mode sweep. Seed the output grad (ones for a scalar), build reverse-topo
        order over ``_parents``, and call each node's ``_backward``.

        TODO: implement the topological walk + seeding. (The per-op ``_backward`` closures do the
        local accumulation; here you only orchestrate the order + seed ``self.grad``.)
        """
        raise NotImplementedError("Tensor.backward")

    def backward(self, grad: Optional[np.ndarray] = None):
        """Reverse-mode sweep. Seed the output grad (ones for a scalar), build reverse-topo
        order over ``_parents``, and call each node's ``_backward``.

        TODO: implement the topological walk + seeding. (The per-op ``_backward`` closures do the
        local accumulation; here you only orchestrate the order + seed ``self.grad``.)
        """
        if grad is None:
            grad = np.ones_like(self.data)
        order = []
        visited = set()
        def topo(t):
            if id(t) not in visited:
                visited.add(id(t))
                for p in t._parents:
                    topo(p)
                order.append(t)
        topo(self)
        self._accum(grad)
        for t in reversed(order):
            if t.grad is not None:
                t._backward()

def add(a, b) -> Tensor:
    """out = a + b (numpy broadcasting). Backward: each input gets _unbroadcast(out.grad)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data + b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        raise NotImplementedError("add backward")  # TODO

    out._backward = _bw
    return out

def add(a, b) -> Tensor:
    """out = a + b (numpy broadcasting). Backward: each input gets _unbroadcast(out.grad)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data + b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        g = out.grad
        if a.requires_grad:
            a._accum(_unbroadcast(g, a.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g, b.shape))

    out._backward = _bw
    return out

def sub(a, b) -> Tensor:
    """out = a - b. Backward: a gets +grad, b gets -grad (both _unbroadcast)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data - b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        raise NotImplementedError("sub backward")  # TODO

    out._backward = _bw
    return out

def sub(a, b) -> Tensor:
    """out = a - b. Backward: a gets +grad, b gets -grad (both _unbroadcast)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data - b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        g = out.grad
        if a.requires_grad:
            a._accum(_unbroadcast(g, a.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(-g, b.shape))

    out._backward = _bw
    return out

def mul(a, b) -> Tensor:
    """out = a * b (elementwise). Backward: a gets grad*b, b gets grad*a (both _unbroadcast)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data * b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        raise NotImplementedError("mul backward")  # TODO

    out._backward = _bw
    return out

def mul(a, b) -> Tensor:
    """out = a * b (elementwise). Backward: a gets grad*b, b gets grad*a (both _unbroadcast)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data * b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        g = out.grad
        if a.requires_grad:
            a._accum(_unbroadcast(g * b.data, a.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g * a.data, b.shape))

    out._backward = _bw
    return out

def div(a, b) -> Tensor:
    """out = a / b (elementwise). Backward: a gets grad/b, b gets -grad*a/b**2 (_unbroadcast)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data / b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        raise NotImplementedError("div backward")  # TODO

    out._backward = _bw
    return out

def div(a, b) -> Tensor:
    """out = a / b (elementwise). Backward: a gets grad/b, b gets -grad*a/b**2 (_unbroadcast)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data / b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        g = out.grad
        if a.requires_grad:
            a._accum(_unbroadcast(g / b.data, a.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(-g * a.data / b.data ** 2, b.shape))

    out._backward = _bw
    return out

def power(a, p: float) -> Tensor:
    """out = a ** p (scalar p). Backward: grad * p * a**(p-1)."""
    a = _ensure(a)
    out = _mk(a.data ** p, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("power backward")  # TODO

    out._backward = _bw
    return out

def power(a, p: float) -> Tensor:
    """out = a ** p (scalar p). Backward: grad * p * a**(p-1)."""
    a = _ensure(a)
    out = _mk(a.data ** p, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * p * a.data ** (p - 1))

    out._backward = _bw
    return out

def maximum(a, b) -> Tensor:
    """elementwise max(a, b). Backward routes grad to whichever input is larger; on EXACT ties,
    split the gradient half-and-half (so finite differences match)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(np.maximum(a.data, b.data), (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        raise NotImplementedError("maximum backward")  # TODO

    out._backward = _bw
    return out

def maximum(a, b) -> Tensor:
    """elementwise max(a, b). Backward routes grad to whichever input is larger; on EXACT ties,
    split the gradient half-and-half (so finite differences match)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(np.maximum(a.data, b.data), (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        g = out.grad
        tie = (a.data == b.data)
        a_wins = (a.data > b.data)
        b_wins = (b.data > a.data)
        if a.requires_grad:
            a._accum(_unbroadcast(g * (a_wins + 0.5 * tie), a.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g * (b_wins + 0.5 * tie), b.shape))

    out._backward = _bw
    return out

def minimum(a, b) -> Tensor:
    """elementwise min(a, b). Backward routes grad to whichever input is smaller (ties split)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(np.minimum(a.data, b.data), (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        raise NotImplementedError("minimum backward")  # TODO

    out._backward = _bw
    return out

def minimum(a, b) -> Tensor:
    """elementwise min(a, b). Backward routes grad to whichever input is smaller (ties split)."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(np.minimum(a.data, b.data), (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        g = out.grad
        tie = (a.data == b.data)
        a_wins = (a.data < b.data)
        b_wins = (b.data < a.data)
        if a.requires_grad:
            a._accum(_unbroadcast(g * (a_wins + 0.5 * tie), a.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g * (b_wins + 0.5 * tie), b.shape))

    out._backward = _bw
    return out

def where(cond, a, b) -> Tensor:
    """np.where(cond, a, b); ``cond`` is a boolean array (NOT differentiated). Backward routes
    grad to ``a`` where cond is True, to ``b`` where False (each _unbroadcast)."""
    a, b = _ensure(a), _ensure(b)
    c = np.asarray(cond, dtype=bool)
    out = _mk(np.where(c, a.data, b.data), (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        raise NotImplementedError("where backward")  # TODO

    out._backward = _bw
    return out

def where(cond, a, b) -> Tensor:
    """np.where(cond, a, b); ``cond`` is a boolean array (NOT differentiated). Backward routes
    grad to ``a`` where cond is True, to ``b`` where False (each _unbroadcast)."""
    a, b = _ensure(a), _ensure(b)
    c = np.asarray(cond, dtype=bool)
    out = _mk(np.where(c, a.data, b.data), (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        g = out.grad
        if a.requires_grad:
            a._accum(_unbroadcast(np.where(c, g, 0), a.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(np.where(c, 0, g), b.shape))

    out._backward = _bw
    return out

def exp(a) -> Tensor:
    """out = exp(a). Backward: grad * exp(a)."""
    a = _ensure(a)
    out = _mk(np.exp(a.data), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("exp backward")  # TODO

    out._backward = _bw
    return out

def exp(a) -> Tensor:
    """out = exp(a). Backward: grad * exp(a)."""
    a = _ensure(a)
    out = _mk(np.exp(a.data), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * out.data)

    out._backward = _bw
    return out

def log(a) -> Tensor:
    """out = log(a). Backward: grad / a."""
    a = _ensure(a)
    out = _mk(np.log(a.data), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("log backward")  # TODO

    out._backward = _bw
    return out

def log(a) -> Tensor:
    """out = log(a). Backward: grad / a."""
    a = _ensure(a)
    out = _mk(np.log(a.data), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad / a.data)

    out._backward = _bw
    return out

def sqrt(a) -> Tensor:
    """out = sqrt(a). Backward: grad * 0.5 / sqrt(a)."""
    a = _ensure(a)
    out = _mk(np.sqrt(a.data), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("sqrt backward")  # TODO

    out._backward = _bw
    return out

def sqrt(a) -> Tensor:
    """out = sqrt(a). Backward: grad * 0.5 / sqrt(a)."""
    a = _ensure(a)
    out = _mk(np.sqrt(a.data), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * 0.5 / out.data)

    out._backward = _bw
    return out

def abs(a) -> Tensor:
    """out = |a|. Backward: grad * sign(a)."""
    a = _ensure(a)
    out = _mk(np.abs(a.data), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("abs backward")  # TODO

    out._backward = _bw
    return out

def abs(a) -> Tensor:
    """out = |a|. Backward: grad * sign(a)."""
    a = _ensure(a)
    out = _mk(np.abs(a.data), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * np.sign(a.data))

    out._backward = _bw
    return out

def clip(a, lo: float, hi: float) -> Tensor:
    """clip to [lo, hi]. Backward passes grad only where lo < a < hi (else 0)."""
    a = _ensure(a)
    out = _mk(np.clip(a.data, lo, hi), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("clip backward")  # TODO

    out._backward = _bw
    return out

def clip(a, lo: float, hi: float) -> Tensor:
    """clip to [lo, hi]. Backward passes grad only where lo < a < hi (else 0)."""
    a = _ensure(a)
    out = _mk(np.clip(a.data, lo, hi), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * ((a.data > lo) & (a.data < hi)))

    out._backward = _bw
    return out

def sum(a, axis=None, keepdims=False) -> Tensor:
    """Sum over ``axis`` (or all). Backward broadcasts out.grad back to a's shape (re-insert a
    summed axis if ``keepdims`` was False)."""
    a = _ensure(a)
    out = _mk(a.data.sum(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("sum backward")  # TODO

    out._backward = _bw
    return out

def _restore_grad(g, axis, keepdims, ndim):
    """Re-insert the reduced axes so g can broadcast back to input shape."""
    if not keepdims:
        if axis is None:
            g = g.reshape([1] * ndim)
        else:
            axes = (axis,) if isinstance(axis, int) else tuple(axis)
            for ax in sorted(axes):
                g = np.expand_dims(g, axis=ax)
    return g


def sum(a, axis=None, keepdims=False) -> Tensor:
    """Sum over ``axis`` (or all). Backward broadcasts out.grad back to a's shape (re-insert a
    summed axis if ``keepdims`` was False)."""
    a = _ensure(a)
    out = _mk(a.data.sum(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            g = _restore_grad(out.grad, axis, keepdims, a.data.ndim)
            a._accum(np.broadcast_to(g, a.shape).copy())

    out._backward = _bw
    return out

def mean(a, axis=None, keepdims=False) -> Tensor:
    """Mean over ``axis`` (or all). Backward = sum-backward scaled by 1/N where N is the number
    of averaged elements."""
    a = _ensure(a)
    out = _mk(a.data.mean(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("mean backward")  # TODO

    out._backward = _bw
    return out

def mean(a, axis=None, keepdims=False) -> Tensor:
    """Mean over ``axis`` (or all). Backward = sum-backward scaled by 1/N where N is the number
    of averaged elements."""
    a = _ensure(a)
    out = _mk(a.data.mean(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            if axis is None:
                N = a.data.size
            else:
                axes = (axis,) if isinstance(axis, int) else tuple(axis)
                N = 1
                for ax in axes:
                    N *= a.data.shape[ax]
            g = _restore_grad(out.grad, axis, keepdims, a.data.ndim)
            a._accum(np.broadcast_to(g / N, a.shape).copy())

    out._backward = _bw
    return out

def max(a, axis=None, keepdims=False) -> Tensor:
    """Reduce-max over ``axis``. Backward routes the upstream grad to the argmax position(s);
    on EXACT ties, split the grad equally among the tied maxima."""
    a = _ensure(a)
    out = _mk(a.data.max(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("max backward")  # TODO

    out._backward = _bw
    return out

def max(a, axis=None, keepdims=False) -> Tensor:
    """Reduce-max over ``axis``. Backward routes the upstream grad to the argmax position(s);
    on EXACT ties, split the grad equally among the tied maxima."""
    a = _ensure(a)
    out = _mk(a.data.max(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            g = _restore_grad(out.grad, axis, keepdims, a.data.ndim)
            max_val = _restore_grad(out.data, axis, keepdims, a.data.ndim)
            mask = (a.data == max_val)
            count = mask.sum(axis=axis, keepdims=True)
            a._accum(g * mask / count)

    out._backward = _bw
    return out

def min(a, axis=None, keepdims=False) -> Tensor:
    """Reduce-min over ``axis``. Backward routes grad to the argmin position(s) (ties split)."""
    a = _ensure(a)
    out = _mk(a.data.min(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("min backward")  # TODO

    out._backward = _bw
    return out

def min(a, axis=None, keepdims=False) -> Tensor:
    """Reduce-min over ``axis``. Backward routes grad to the argmin position(s) (ties split)."""
    a = _ensure(a)
    out = _mk(a.data.min(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            g = _restore_grad(out.grad, axis, keepdims, a.data.ndim)
            min_val = _restore_grad(out.data, axis, keepdims, a.data.ndim)
            mask = (a.data == min_val)
            count = mask.sum(axis=axis, keepdims=True)
            a._accum(g * mask / count)

    out._backward = _bw
    return out

def concat(tensors, axis=0) -> Tensor:
    """Concatenate Tensors along ``axis``. Backward splits out.grad back to each input (each input
    receives the slice of out.grad covering its span along ``axis``)."""
    ts = [_ensure(t) for t in tensors]
    out = _mk(np.concatenate([t.data for t in ts], axis=axis), tuple(ts),
              any(t.requires_grad for t in ts))

    def _bw():
        raise NotImplementedError("concat backward")  # TODO

    out._backward = _bw
    return out

def concat(tensors, axis=0) -> Tensor:
    """Concatenate Tensors along ``axis``. Backward splits out.grad back to each input (each input
    receives the slice of out.grad covering its span along ``axis``)."""
    ts = [_ensure(t) for t in tensors]
    out = _mk(np.concatenate([t.data for t in ts], axis=axis), tuple(ts),
              any(t.requires_grad for t in ts))

    def _bw():
        g = out.grad
        sizes = [t.data.shape[axis] for t in ts]
        splits = np.split(g, np.cumsum(sizes[:-1]), axis=axis)
        for t, sg in zip(ts, splits):
            if t.requires_grad:
                t._accum(sg)

    out._backward = _bw
    return out

def stack(tensors, axis=0) -> Tensor:
    """Stack Tensors along a NEW ``axis``. Backward indexes out.grad along that axis per input."""
    ts = [_ensure(t) for t in tensors]
    out = _mk(np.stack([t.data for t in ts], axis=axis), tuple(ts),
              any(t.requires_grad for t in ts))

    def _bw():
        raise NotImplementedError("stack backward")  # TODO

    out._backward = _bw
    return out

def stack(tensors, axis=0) -> Tensor:
    """Stack Tensors along a NEW ``axis``. Backward indexes out.grad along that axis per input."""
    ts = [_ensure(t) for t in tensors]
    out = _mk(np.stack([t.data for t in ts], axis=axis), tuple(ts),
              any(t.requires_grad for t in ts))

    def _bw():
        g = out.grad
        for i, t in enumerate(ts):
            if t.requires_grad:
                idx = [slice(None)] * g.ndim
                idx[axis] = i
                t._accum(g[tuple(idx)])

    out._backward = _bw
    return out

def transpose(a, axes=None) -> Tensor:
    """np.transpose(a, axes). Backward: transpose out.grad by the inverse permutation."""
    a = _ensure(a)
    out = _mk(np.transpose(a.data, axes), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("transpose backward")  # TODO

    out._backward = _bw
    return out

def transpose(a, axes=None) -> Tensor:
    """np.transpose(a, axes). Backward: transpose out.grad by the inverse permutation."""
    a = _ensure(a)
    out = _mk(np.transpose(a.data, axes), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            if axes is None:
                inv_axes = None
            else:
                inv_axes = np.argsort(axes)
            a._accum(np.transpose(out.grad, inv_axes))

    out._backward = _bw
    return out

def reshape(a, shape) -> Tensor:
    """Reshape. Backward: reshape out.grad back to a's original shape."""
    a = _ensure(a)
    out = _mk(a.data.reshape(shape), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("reshape backward")  # TODO

    out._backward = _bw
    return out

def reshape(a, shape) -> Tensor:
    """Reshape. Backward: reshape out.grad back to a's original shape."""
    a = _ensure(a)
    out = _mk(a.data.reshape(shape), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad.reshape(a.shape))

    out._backward = _bw
    return out

def getitem(a, idx) -> Tensor:
    """Index/slice a tensor (``a[idx]``). Backward scatter-adds out.grad into a zeros-like(a) at
    ``idx``."""
    a = _ensure(a)
    out = _mk(a.data[idx], (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("getitem backward")  # TODO

    out._backward = _bw
    return out

def getitem(a, idx) -> Tensor:
    """Index/slice a tensor (``a[idx]``). Backward scatter-adds out.grad into a zeros-like(a) at
    ``idx``."""
    a = _ensure(a)
    out = _mk(a.data[idx], (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            np.add.at(da, idx, out.grad)
            a._accum(da)

    out._backward = _bw
    return out

def matmul(a, b) -> Tensor:
    """Batched matmul (a @ b) with numpy broadcasting over leading dims. Backward:
    da = grad @ b^T, db = a^T @ grad (swap the LAST TWO axes for the transposes), each
    _unbroadcast back to the input's shape so broadcasted batch dims are reduced."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data @ b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        raise NotImplementedError("matmul backward")  # TODO

    out._backward = _bw
    return out

def matmul(a, b) -> Tensor:
    """Batched matmul (a @ b) with numpy broadcasting over leading dims. Backward:
    da = grad @ b^T, db = a^T @ grad (swap the LAST TWO axes for the transposes), each
    _unbroadcast back to the input's shape so broadcasted batch dims are reduced."""
    a, b = _ensure(a), _ensure(b)
    out = _mk(a.data @ b.data, (a, b), a.requires_grad or b.requires_grad)

    def _bw():
        g = out.grad
        if a.requires_grad:
            da = g @ np.swapaxes(b.data, -1, -2)
            a._accum(_unbroadcast(da, a.shape))
        if b.requires_grad:
            db = np.swapaxes(a.data, -1, -2) @ g
            b._accum(_unbroadcast(db, b.shape))

    out._backward = _bw
    return out

def relu(a) -> Tensor:
    """max(a, 0). Backward: grad * (a > 0)."""
    a = _ensure(a)
    out = _mk(np.maximum(a.data, 0.0), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("relu backward")  # TODO

    out._backward = _bw
    return out

def relu(a) -> Tensor:
    """max(a, 0). Backward: grad * (a > 0)."""
    a = _ensure(a)
    out = _mk(np.maximum(a.data, 0.0), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * (a.data > 0))

    out._backward = _bw
    return out

def leaky_relu(a, slope: float = 0.01) -> Tensor:
    """out = a where a>0 else slope*a. Backward: grad where a>0, slope*grad where a<=0."""
    a = _ensure(a)
    out = _mk(np.where(a.data > 0.0, a.data, slope * a.data), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("leaky_relu backward")  # TODO

    out._backward = _bw
    return out

def leaky_relu(a, slope: float = 0.01) -> Tensor:
    """out = a where a>0 else slope*a. Backward: grad where a>0, slope*grad where a<=0."""
    a = _ensure(a)
    out = _mk(np.where(a.data > 0.0, a.data, slope * a.data), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * np.where(a.data > 0, 1.0, slope))

    out._backward = _bw
    return out

def sigmoid(a) -> Tensor:
    """out = 1/(1+exp(-a)). Backward: grad * out * (1 - out)."""
    a = _ensure(a)
    with np.errstate(over="ignore"):
        s = 1.0 / (1.0 + np.exp(-a.data))
    out = _mk(s, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("sigmoid backward")  # TODO

    out._backward = _bw
    return out

def sigmoid(a) -> Tensor:
    """out = 1/(1+exp(-a)). Backward: grad * out * (1 - out)."""
    a = _ensure(a)
    with np.errstate(over="ignore"):
        s = 1.0 / (1.0 + np.exp(-a.data))
    out = _mk(s, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * out.data * (1.0 - out.data))

    out._backward = _bw
    return out

def tanh(a) -> Tensor:
    """out = tanh(a). Backward: grad * (1 - tanh(a)**2)."""
    a = _ensure(a)
    out = _mk(np.tanh(a.data), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("tanh backward")  # TODO

    out._backward = _bw
    return out

def tanh(a) -> Tensor:
    """out = tanh(a). Backward: grad * (1 - tanh(a)**2)."""
    a = _ensure(a)
    out = _mk(np.tanh(a.data), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * (1.0 - out.data ** 2))

    out._backward = _bw
    return out

def gelu(a) -> Tensor:
    """Exact-erf GELU: out = a * 0.5 * (1 + erf(a / sqrt(2))).  (``_erf`` is provided above.)
    Backward: grad * (cdf + a * pdf), cdf = 0.5*(1+erf(a/sqrt2)), pdf = exp(-a^2/2)/sqrt(2*pi)."""
    a = _ensure(a)
    cdf = 0.5 * (1.0 + _erf(a.data * _INV_SQRT2))
    out = _mk(a.data * cdf, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("gelu backward")  # TODO

    out._backward = _bw
    return out

def gelu(a) -> Tensor:
    """Exact-erf GELU: out = a * 0.5 * (1 + erf(a / sqrt(2))).  (``_erf`` is provided above.)
    Backward: grad * (cdf + a * pdf), cdf = 0.5*(1+erf(a/sqrt2)), pdf = exp(-a^2/2)/sqrt(2*pi)."""
    a = _ensure(a)
    cdf = 0.5 * (1.0 + _erf(a.data * _INV_SQRT2))
    out = _mk(a.data * cdf, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            pdf = np.exp(-0.5 * a.data ** 2) / np.sqrt(2.0 * np.pi)
            a._accum(out.grad * (cdf + a.data * pdf))

    out._backward = _bw
    return out

def softmax(a, axis=-1) -> Tensor:
    """Numerically-stable softmax over ``axis`` (subtract the max before exp).
    Backward (Jacobian-vector product, no full Jacobian):
        s = softmax(a);  da = s * (grad - sum(grad * s, axis, keepdims))."""
    a = _ensure(a)
    z = a.data - a.data.max(axis=axis, keepdims=True)
    with np.errstate(under="ignore"):
        e = np.exp(z)
        s = e / e.sum(axis=axis, keepdims=True)
    out = _mk(s, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("softmax backward")  # TODO

    out._backward = _bw
    return out

def softmax(a, axis=-1) -> Tensor:
    """Numerically-stable softmax over ``axis`` (subtract the max before exp).
    Backward (Jacobian-vector product, no full Jacobian):
        s = softmax(a);  da = s * (grad - sum(grad * s, axis, keepdims))."""
    a = _ensure(a)
    z = a.data - a.data.max(axis=axis, keepdims=True)
    with np.errstate(under="ignore"):
        e = np.exp(z)
        s = e / e.sum(axis=axis, keepdims=True)
    out = _mk(s, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            g = out.grad
            a._accum(out.data * (g - (g * out.data).sum(axis=axis, keepdims=True)))

    out._backward = _bw
    return out

def log_softmax(a, axis=-1) -> Tensor:
    """Stable log-softmax: z - logsumexp(z). Backward: grad - softmax * sum(grad, axis, keep)."""
    a = _ensure(a)
    z = a.data - a.data.max(axis=axis, keepdims=True)
    with np.errstate(under="ignore"):
        ls = z - np.log(np.exp(z).sum(axis=axis, keepdims=True))
    out = _mk(ls, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("log_softmax backward")  # TODO

    out._backward = _bw
    return out

def log_softmax(a, axis=-1) -> Tensor:
    """Stable log-softmax: z - logsumexp(z). Backward: grad - softmax * sum(grad, axis, keep)."""
    a = _ensure(a)
    z = a.data - a.data.max(axis=axis, keepdims=True)
    with np.errstate(under="ignore"):
        ls = z - np.log(np.exp(z).sum(axis=axis, keepdims=True))
    out = _mk(ls, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            g = out.grad
            a._accum(g - np.exp(out.data) * g.sum(axis=axis, keepdims=True))

    out._backward = _bw
    return out

def cross_entropy(logits, targets) -> Tensor:
    """Mean cross-entropy over a batch. ``logits``: (N, C). ``targets``: int (N,). Returns a
    SCALAR Tensor = mean over N of -log_softmax(logits)[i, targets[i]]. Stable (log-sum-exp)
    forward is provided. Backward into logits: (softmax - onehot(targets)) / N, times out.grad.
    """
    logits = _ensure(logits)
    x = logits.data
    z = x - x.max(axis=-1, keepdims=True)
    e = np.exp(z)
    logp = z - np.log(e.sum(axis=-1, keepdims=True))
    t = np.asarray(targets).astype(np.int64).reshape(-1)
    n = x.shape[0]
    out = _mk(-logp[np.arange(n), t].mean(), (logits,), logits.requires_grad)

    def _bw():
        raise NotImplementedError("cross_entropy backward")  # TODO

    out._backward = _bw
    return out

def cross_entropy(logits, targets) -> Tensor:
    """Mean cross-entropy over a batch. ``logits``: (N, C). ``targets``: int (N,). Returns a
    SCALAR Tensor = mean over N of -log_softmax(logits)[i, targets[i]]. Stable (log-sum-exp)
    forward is provided. Backward into logits: (softmax - onehot(targets)) / N, times out.grad.
    """
    logits = _ensure(logits)
    x = logits.data
    z = x - x.max(axis=-1, keepdims=True)
    e = np.exp(z)
    logp = z - np.log(e.sum(axis=-1, keepdims=True))
    t = np.asarray(targets).astype(np.int64).reshape(-1)
    n = x.shape[0]
    out = _mk(-logp[np.arange(n), t].mean(), (logits,), logits.requires_grad)

    def _bw():
        if logits.requires_grad:
            sm = e / e.sum(axis=-1, keepdims=True)
            onehot = np.zeros_like(sm)
            onehot[np.arange(n), t] = 1.0
            dlogits = (sm - onehot) / n * out.grad
            logits._accum(dlogits)

    out._backward = _bw
    return out

def mse_loss(pred, target) -> Tensor:
    """Mean squared error over ALL elements: mean((pred - target)^2). ``target`` is a constant
    array. Backward into pred: (2/N) * (pred - target) * out.grad, N = pred.size."""
    pred = _ensure(pred)
    tgt = np.asarray(target, dtype=np.float64)
    out = _mk(float(np.mean((pred.data - tgt) ** 2)), (pred,), pred.requires_grad)

    def _bw():
        raise NotImplementedError("mse_loss backward")  # TODO

    out._backward = _bw
    return out

def mse_loss(pred, target) -> Tensor:
    """Mean squared error over ALL elements: mean((pred - target)^2). ``target`` is a constant
    array. Backward into pred: (2/N) * (pred - target) * out.grad, N = pred.size."""
    pred = _ensure(pred)
    tgt = np.asarray(target, dtype=np.float64)
    out = _mk(float(np.mean((pred.data - tgt) ** 2)), (pred,), pred.requires_grad)

    def _bw():
        if pred.requires_grad:
            pred._accum(2.0 * (pred.data - tgt) / pred.data.size * out.grad)

    out._backward = _bw
    return out

def layernorm(x, gamma, beta, eps: float = 1e-5) -> Tensor:
    """Layer normalization over the LAST axis. x:(...,D); gamma,beta:(D,). POPULATION variance.
        mu=mean(x,-1); xhat=(x-mu)/sqrt(var+eps); out = gamma*xhat + beta.
    Backward (sums over the last axis): dgamma=sum(grad*xhat over leading axes);
      dbeta=sum(grad over leading axes); gx=grad*gamma;
      dx = inv/D * (D*gx - sum(gx,-1) - xhat*sum(gx*xhat,-1)).
    The forward (mu/var/inv/xhat) is provided; implement the backward."""
    x, gamma, beta = _ensure(x), _ensure(gamma), _ensure(beta)
    xd = x.data
    D = xd.shape[-1]
    mu = xd.mean(axis=-1, keepdims=True)
    xc = xd - mu
    var = (xc * xc).mean(axis=-1, keepdims=True)
    inv = 1.0 / np.sqrt(var + eps)
    xhat = xc * inv
    out = _mk(gamma.data * xhat + beta.data, (x, gamma, beta),
              x.requires_grad or gamma.requires_grad or beta.requires_grad)

    def _bw():
        raise NotImplementedError("layernorm backward")  # TODO

    out._backward = _bw
    return out

def layernorm(x, gamma, beta, eps: float = 1e-5) -> Tensor:
    """Layer normalization over the LAST axis. x:(...,D); gamma,beta:(D,). POPULATION variance.
        mu=mean(x,-1); xhat=(x-mu)/sqrt(var+eps); out = gamma*xhat + beta.
    Backward (sums over the last axis): dgamma=sum(grad*xhat over leading axes);
      dbeta=sum(grad over leading axes); gx=grad*gamma;
      dx = inv/D * (D*gx - sum(gx,-1) - xhat*sum(gx*xhat,-1)).
    The forward (mu/var/inv/xhat) is provided; implement the backward."""
    x, gamma, beta = _ensure(x), _ensure(gamma), _ensure(beta)
    xd = x.data
    D = xd.shape[-1]
    mu = xd.mean(axis=-1, keepdims=True)
    xc = xd - mu
    var = (xc * xc).mean(axis=-1, keepdims=True)
    inv = 1.0 / np.sqrt(var + eps)
    xhat = xc * inv
    out = _mk(gamma.data * xhat + beta.data, (x, gamma, beta),
              x.requires_grad or gamma.requires_grad or beta.requires_grad)

    def _bw():
        g = out.grad
        if gamma.requires_grad:
            # sum over all leading axes (all except last)
            leading = tuple(range(g.ndim - 1))
            gamma._accum((g * xhat).sum(axis=leading))
        if beta.requires_grad:
            leading = tuple(range(g.ndim - 1))
            beta._accum(g.sum(axis=leading))
        if x.requires_grad:
            gx = g * gamma.data
            dx = inv / D * (D * gx - gx.sum(axis=-1, keepdims=True) - xhat * (gx * xhat).sum(axis=-1, keepdims=True))
            x._accum(dx)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("conv2d backward")  # TODO (dW, dbias, dx via _col2im)

    out._backward = _bw
    return out

    def _bw():
        dout_mat = out.grad.reshape(N, Cout, OH * OW)
        if weight.requires_grad:
            dW = np.einsum("nop,ncp->oc", dout_mat, cols).reshape(Cout, Cin, KH, KW)
            weight._accum(dW)
        if has_bias and bias.requires_grad:
            dbias = dout_mat.sum(axis=(0, 2))
            bias._accum(dbias)
        if x.requires_grad:
            dcols = np.einsum("oc,nop->ncp", Wm, dout_mat)
            dxp = _col2im(dcols, xp.shape, KH, KW, stride, OH, OW)
            dx = dxp[:, :, pad:-pad, pad:-pad] if pad > 0 else dxp
            x._accum(dx)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("avgpool2d backward")  # TODO

    out._backward = _bw
    return out


def maxpool2d(x, k) -> Tensor:

    def _bw():
        if x.requires_grad:
            dx = np.zeros_like(x.data)
            dx[:, :, :OH * k, :OW * k] = np.repeat(np.repeat(out.grad, k, axis=2), k, axis=3) / (k * k)
            x._accum(dx)

    out._backward = _bw
    return out


def maxpool2d(x, k) -> Tensor:

    def _bw():
        raise NotImplementedError("maxpool2d backward")  # TODO

    out._backward = _bw
    return out


def batchnorm2d

    def _bw():
        if x.requires_grad:
            dx = np.zeros_like(x.data)
            # xr shape: (N, C, OH, k, OW, k) -> transpose to (N, C, OH, OW, k, k)
            xr_t = xr.transpose(0, 1, 2, 4, 3, 5).reshape(N, C, OH, OW, k * k)
            af = xr_t.argmax(axis=-1)
            ai = af // k
            aj = af % k
            hi = np.arange(OH)[None, None, :, None] * k + ai
            wj = np.arange(OW)[None, None, None, :] * k + aj
            n_idx = np.arange(N)[:, None, None, None]
            c_idx = np.arange(C)[None, :, None, None]
            np.add.at(dx, (n_idx, c_idx, hi, wj), out.grad)
            x._accum(dx)

    out._backward = _bw
    return out


def batchnorm2d

    def _bw():
        raise NotImplementedError("batchnorm2d backward")  # TODO

    out._backward = _bw
    return out


# ============================================================ LSQ fake-quant (DOMAIN; the crux)

    def _bw():
        g = out.grad
        if gamma.requires_grad:
            gamma._accum((g * xhat).sum(axis=(0, 2, 3)))
        if beta.requires_grad:
            beta._accum(g.sum(axis=(0, 2, 3)))
        if x.requires_grad:
            if training:
                gx = g * g_
                dx = inv / M * (M * gx - gx.sum(axis=(0, 2, 3), keepdims=True) - xhat * (gx * xhat).sum(axis=(0, 2, 3), keepdims=True))
            else:
                dx = g * g_ * inv
            x._accum(dx)

    out._backward = _bw
    return out


# ============================================================ LSQ fake-quant (DOMAIN; the crux)

    def _bw():
        raise NotImplementedError("fake_quant backward (STE + LSQ scale gradient)")  # TODO

    out._backward = _bw
    return out

    def _bw():
        g = out.grad
        ste_mask = (r > Qn) & (r < Qp)
        if v.requires_grad:
            v._accum(g * ste_mask)
        if scale.requires_grad:
            below = r < Qn
            above = r > Qp
            d = np.where(below, Qn, np.where(above, Qp, v_hat - r))
            ds = float((g * d).sum()) * grad_scale
            scale._accum(np.array([ds]))

    out._backward = _bw
    return out

def square(a) -> Tensor:
    """out = a*a. Backward: grad * 2*a."""
    a = _ensure(a)
    out = _mk(a.data * a.data, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("square backward")  # TODO

    out._backward = _bw
    return out

def square(a) -> Tensor:
    """out = a*a. Backward: grad * 2*a."""
    a = _ensure(a)
    out = _mk(a.data * a.data, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * 2.0 * a.data)

    out._backward = _bw
    return out

def rsqrt(a) -> Tensor:
    """out = 1/sqrt(a). Backward: grad * -0.5 * a**(-1.5)."""
    a = _ensure(a)
    out = _mk(1.0 / np.sqrt(a.data), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("rsqrt backward")  # TODO

    out._backward = _bw
    return out

def rsqrt(a) -> Tensor:
    """out = 1/sqrt(a). Backward: grad * -0.5 * a**(-1.5)."""
    a = _ensure(a)
    out = _mk(1.0 / np.sqrt(a.data), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * -0.5 * a.data ** (-1.5))

    out._backward = _bw
    return out

def reciprocal(a) -> Tensor:
    """out = 1/a. Backward: grad * -1/a^2."""
    a = _ensure(a)
    out = _mk(1.0 / a.data, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("reciprocal backward")  # TODO

    out._backward = _bw
    return out

def reciprocal(a) -> Tensor:
    """out = 1/a. Backward: grad * -1/a^2."""
    a = _ensure(a)
    out = _mk(1.0 / a.data, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * -1.0 / a.data ** 2)

    out._backward = _bw
    return out

def var(a, axis=None, keepdims=False) -> Tensor:
    """Population variance over ``axis``: mean((a-mu)^2). Backward: grad * (2/N) * (a - mu)."""
    a = _ensure(a)
    mu = a.data.mean(axis=axis, keepdims=True)
    xc = a.data - mu
    out = _mk((xc * xc).mean(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("var backward")  # TODO

    out._backward = _bw
    return out

def var(a, axis=None, keepdims=False) -> Tensor:
    """Population variance over ``axis``: mean((a-mu)^2). Backward: grad * (2/N) * (a - mu)."""
    a = _ensure(a)
    mu = a.data.mean(axis=axis, keepdims=True)
    xc = a.data - mu
    out = _mk((xc * xc).mean(axis=axis, keepdims=keepdims), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            if axis is None:
                N = a.data.size
            else:
                axes = (axis,) if isinstance(axis, int) else tuple(axis)
                N = 1
                for ax in axes:
                    N *= a.data.shape[ax]
            g_e = _restore_grad(out.grad, axis, keepdims, a.data.ndim)
            a._accum(np.broadcast_to(g_e, a.shape).copy() * (2.0 / N) * xc)

    out._backward = _bw
    return out

def std(a, axis=None, keepdims=False, eps: float = 0.0) -> Tensor:
    """Population standard deviation sqrt(var + eps). Backward: grad * (a-mu)/(N*std)."""
    a = _ensure(a)
    mu = a.data.mean(axis=axis, keepdims=True)
    xc = a.data - mu
    v = (xc * xc).mean(axis=axis, keepdims=True)
    sd = np.sqrt(v + eps)
    out_data = sd if keepdims else (np.squeeze(sd, axis=axis) if axis is not None else float(sd.reshape(())))
    out = _mk(out_data, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("std backward")  # TODO

    out._backward = _bw
    return out

def std(a, axis=None, keepdims=False, eps: float = 0.0) -> Tensor:
    """Population standard deviation sqrt(var + eps). Backward: grad * (a-mu)/(N*std)."""
    a = _ensure(a)
    mu = a.data.mean(axis=axis, keepdims=True)
    xc = a.data - mu
    v = (xc * xc).mean(axis=axis, keepdims=True)
    sd = np.sqrt(v + eps)
    out_data = sd if keepdims else (np.squeeze(sd, axis=axis) if axis is not None else float(sd.reshape(())))
    out = _mk(out_data, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            if axis is None:
                N = a.data.size
            else:
                axes = (axis,) if isinstance(axis, int) else tuple(axis)
                N = 1
                for ax in axes:
                    N *= a.data.shape[ax]
            g_e = _restore_grad(out.grad, axis, keepdims, a.data.ndim)
            # sd has keepdims=True shape already
            a._accum(np.broadcast_to(g_e, a.shape).copy() * xc / (N * sd))

    out._backward = _bw
    return out

def gather(a, idx, axis: int) -> Tensor:
    """Gather along ``axis`` (numpy take_along_axis). Backward scatter-adds out.grad back at idx."""
    a = _ensure(a)
    idx = np.asarray(idx)
    out = _mk(np.take_along_axis(a.data, idx, axis=axis), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("gather backward")  # TODO

    out._backward = _bw
    return out

def gather(a, idx, axis: int) -> Tensor:
    """Gather along ``axis`` (numpy take_along_axis). Backward scatter-adds out.grad back at idx."""
    a = _ensure(a)
    idx = np.asarray(idx)
    out = _mk(np.take_along_axis(a.data, idx, axis=axis), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            np.add.at(da, tuple(np.indices(idx.shape)[:axis]) + (idx,) + tuple(np.indices(idx.shape)[axis+1:]), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            np.add.at(da, tuple(np.indices(idx.shape)[:axis]) + (idx,) + tuple(np.indices(idx.shape)[axis+1:]), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            np.add.at(da, np.ix_(*[np.arange(s) for s in idx.shape[:axis]]) + (idx,) + np.ix_(*[np.arange(s) for s in idx.shape[axis+1:]]) if idx.ndim > 1 else (idx,), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            np.add.at(da, np.ix_(*[np.arange(s) for s in idx.shape[:axis]]) + (idx,) + np.ix_(*[np.arange(s) for s in idx.shape[axis+1:]]) if idx.ndim > 1 else (idx,), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            # Build full index tuple so np.add.at scatters along 'axis'
            full_idx = []
            for d in range(a.data.ndim):
                if d == axis:
                    full_idx.append(idx)
                else:
                    shape = [1] * idx.ndim
                    shape[d if d < axis else d] = idx.shape[d if d < axis else d]
                    rng = np.arange(idx.shape[d if d < axis else d])
                    # broadcast-compatible arange
                    bshape = [1] * idx.ndim
                    bshape[d] = idx.shape[d]
                    full_idx.append(np.arange(idx.shape[d]).reshape(bshape))
            np.add.at(da, tuple(full_idx), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            # Build full index tuple so np.add.at scatters along 'axis'
            full_idx = []
            for d in range(a.data.ndim):
                if d == axis:
                    full_idx.append(idx)
                else:
                    shape = [1] * idx.ndim
                    shape[d if d < axis else d] = idx.shape[d if d < axis else d]
                    rng = np.arange(idx.shape[d if d < axis else d])
                    # broadcast-compatible arange
                    bshape = [1] * idx.ndim
                    bshape[d] = idx.shape[d]
                    full_idx.append(np.arange(idx.shape[d]).reshape(bshape))
            np.add.at(da, tuple(full_idx), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            np.add.at(da, _gather_idx(idx, axis, a.data.ndim), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

def gather(a, idx, axis: int) -> Tensor:
    """Gather along ``axis`` (numpy take_along_axis). Backward scatter-adds out.grad back at idx."""
    a = _ensure(a)
    idx = np.asarray(idx)
    out = _mk(np.take_along_axis(a.data, idx, axis=axis), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            np.add.at(da, _gather_idx(idx, axis, a.data.ndim), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

def gather(a, idx, axis: int) -> Tensor:
    """Gather along ``axis`` (numpy take_along_axis). Backward scatter-adds out.grad back at idx."""
    a = _ensure(a)
    idx = np.asarray(idx)
    out = _mk(np.take_along_axis(a.data, idx, axis=axis), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            da = np.zeros_like(a.data)
            # np.put_along_axis does scatter-add equivalent but replaces; use add.at
            # Build broadcast index tuple compatible with np.add.at
            ndim = a.data.ndim
            full_idx = []
            for d in range(ndim):
                if d == axis:
                    full_idx.append(idx)
                else:
                    shape = [1] * ndim
                    shape[d] = idx.shape[d]
                    full_idx.append(np.arange(idx.shape[d]).reshape(shape))
            np.add.at(da, tuple(full_idx), out.grad)
            a._accum(da)

    out._backward = _bw
    return out

def pad2d(x, pad: int) -> Tensor:
    """Zero-pad the last two (H,W) axes of an NCHW tensor by ``pad`` on each side. Backward crops
    out.grad back to the original H,W. (``_pad_nchw`` is provided above.)"""
    x = _ensure(x)
    out = _mk(_pad_nchw(x.data, pad), (x,), x.requires_grad)

    def _bw():
        raise NotImplementedError("pad2d backward")  # TODO

    out._backward = _bw
    return out

def pad2d(x, pad: int) -> Tensor:
    """Zero-pad the last two (H,W) axes of an NCHW tensor by ``pad`` on each side. Backward crops
    out.grad back to the original H,W. (``_pad_nchw`` is provided above.)"""
    x = _ensure(x)
    out = _mk(_pad_nchw(x.data, pad), (x,), x.requires_grad)

    def _bw():
        if x.requires_grad:
            if pad > 0:
                x._accum(out.grad[:, :, pad:-pad, pad:-pad])
            else:
                x._accum(out.grad)

    out._backward = _bw
    return out

def softplus(a, beta: float = 1.0) -> Tensor:
    """out = (1/beta)*log(1+exp(beta*a)) (stable form provided). Backward: grad * sigmoid(beta*a)."""
    a = _ensure(a)
    bx = beta * a.data
    out = _mk((np.maximum(bx, 0.0) + np.log1p(np.exp(-np.abs(bx)))) / beta, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("softplus backward")  # TODO

    out._backward = _bw
    return out

def softplus(a, beta: float = 1.0) -> Tensor:
    """out = (1/beta)*log(1+exp(beta*a)) (stable form provided). Backward: grad * sigmoid(beta*a)."""
    a = _ensure(a)
    bx = beta * a.data
    out = _mk((np.maximum(bx, 0.0) + np.log1p(np.exp(-np.abs(bx)))) / beta, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            with np.errstate(over="ignore"):
                sig = 1.0 / (1.0 + np.exp(-bx))
            a._accum(out.grad * sig)

    out._backward = _bw
    return out

def silu(a) -> Tensor:
    """SiLU / swish: out = a*sigmoid(a). Backward: grad*(sig + a*sig*(1-sig)), sig=sigmoid(a)."""
    a = _ensure(a)
    with np.errstate(over="ignore"):
        sig = 1.0 / (1.0 + np.exp(-a.data))
    out = _mk(a.data * sig, (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("silu backward")  # TODO

    out._backward = _bw
    return out

def silu(a) -> Tensor:
    """SiLU / swish: out = a*sigmoid(a). Backward: grad*(sig + a*sig*(1-sig)), sig=sigmoid(a)."""
    a = _ensure(a)
    with np.errstate(over="ignore"):
        sig = 1.0 / (1.0 + np.exp(-a.data))
    out = _mk(a.data * sig, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * (sig + a.data * sig * (1.0 - sig)))

    out._backward = _bw
    return out

def mish(a) -> Tensor:
    """Mish: out = a*tanh(softplus(a)). Backward via the chain rule (softplus'(x)=sigmoid(x))."""
    a = _ensure(a)
    x = a.data
    sp = np.maximum(x, 0.0) + np.log1p(np.exp(-np.abs(x)))
    out = _mk(x * np.tanh(sp), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("mish backward")  # TODO

    out._backward = _bw
    return out

def mish(a) -> Tensor:
    """Mish: out = a*tanh(softplus(a)). Backward via the chain rule (softplus'(x)=sigmoid(x))."""
    a = _ensure(a)
    x = a.data
    sp = np.maximum(x, 0.0) + np.log1p(np.exp(-np.abs(x)))
    th_sp = np.tanh(sp)
    out = _mk(x * th_sp, (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            with np.errstate(over="ignore"):
                sig = 1.0 / (1.0 + np.exp(-x))
            # d/dx (x * tanh(sp)) = tanh(sp) + x * (1 - tanh(sp)^2) * sigmoid(x)
            a._accum(out.grad * (th_sp + x * (1.0 - th_sp ** 2) * sig))

    out._backward = _bw
    return out

def elu(a, alpha: float = 1.0) -> Tensor:
    """ELU: out = a where a>0 else alpha*(exp(a)-1). Backward: grad where a>0 else grad*alpha*exp(a)."""
    a = _ensure(a)
    x = a.data
    out = _mk(np.where(x > 0.0, x, alpha * (np.exp(np.minimum(x, 0.0)) - 1.0)), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("elu backward")  # TODO

    out._backward = _bw
    return out

def elu(a, alpha: float = 1.0) -> Tensor:
    """ELU: out = a where a>0 else alpha*(exp(a)-1). Backward: grad where a>0 else grad*alpha*exp(a)."""
    a = _ensure(a)
    x = a.data
    out = _mk(np.where(x > 0.0, x, alpha * (np.exp(np.minimum(x, 0.0)) - 1.0)), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * np.where(x > 0.0, 1.0, alpha * np.exp(np.minimum(x, 0.0))))

    out._backward = _bw
    return out

def hardtanh(a, lo: float = -1.0, hi: float = 1.0) -> Tensor:
    """Clamp to [lo,hi]. Backward passes grad only where lo < a < hi (else 0)."""
    a = _ensure(a)
    out = _mk(np.clip(a.data, lo, hi), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("hardtanh backward")  # TODO

    out._backward = _bw
    return out

def hardtanh(a, lo: float = -1.0, hi: float = 1.0) -> Tensor:
    """Clamp to [lo,hi]. Backward passes grad only where lo < a < hi (else 0)."""
    a = _ensure(a)
    out = _mk(np.clip(a.data, lo, hi), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad * ((a.data > lo) & (a.data < hi)))

    out._backward = _bw
    return out

def hardsigmoid(a) -> Tensor:
    """out = clip(a/6 + 0.5, 0, 1). Backward: grad/6 where 0 < a/6+0.5 < 1 else 0."""
    a = _ensure(a)
    z = a.data / 6.0 + 0.5
    out = _mk(np.clip(z, 0.0, 1.0), (a,), a.requires_grad)

    def _bw():
        raise NotImplementedError("hardsigmoid backward")  # TODO

    out._backward = _bw
    return out

def hardsigmoid(a) -> Tensor:
    """out = clip(a/6 + 0.5, 0, 1). Backward: grad/6 where 0 < a/6+0.5 < 1 else 0."""
    a = _ensure(a)
    z = a.data / 6.0 + 0.5
    out = _mk(np.clip(z, 0.0, 1.0), (a,), a.requires_grad)

    def _bw():
        if a.requires_grad:
            a._accum(out.grad / 6.0 * ((z > 0.0) & (z < 1.0)))

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("groupnorm backward")  # TODO

    out._backward = _bw
    return out


# ============================================================ EXTRA quantization (per-channel + affine)

    def _bw():
        g = out.grad
        if gamma.requires_grad:
            gamma._accum((g * xhat).sum(axis=(0, 2, 3)))
        if beta.requires_grad:
            beta._accum(g.sum(axis=(0, 2, 3)))
        if x.requires_grad:
            M = cg * H * W
            gx = (g * gamma.data.reshape(1, C, 1, 1)).reshape(N, G, M)
            inv_g = inv  # shape (N, G, 1)
            xhat_r = xhat.reshape(N, G, M)
            dx = inv_g / M * (M * gx - gx.sum(axis=-1, keepdims=True) - xhat_r * (gx * xhat_r).sum(axis=-1, keepdims=True))
            x._accum(dx.reshape(N, C, H, W))

    out._backward = _bw
    return out


# ============================================================ EXTRA quantization (per-channel + affine)

    def _bw():
        raise NotImplementedError("fake_quant_per_channel backward (STE + per-channel scale grad)")  # TODO

    out._backward = _bw
    return out

    def _bw():
        g = out.grad
        ste_mask = (r > Qn) & (r < Qp)
        if v.requires_grad:
            v._accum(g * ste_mask)
        if scale.requires_grad:
            d = np.where(r < Qn, Qn, np.where(r > Qp, Qp, v_hat - r))
            axes_to_sum = tuple(i for i in range(v.data.ndim) if i != axis)
            scale._accum((g * d).sum(axis=axes_to_sum) * grad_scale)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("fake_quant_affine backward (STE + scale grad on shifted grid)")  # TODO

    out._backward = _bw
    return out

    def _bw():
        g = out.grad
        ste_mask = (r > Qn) & (r < Qp)
        if v.requires_grad:
            v._accum(g * ste_mask)
        if scale.requires_grad:
            # d = (q - z) - (r - z) in-range = q - r in-range
            # but: (Qn - z) below, (Qp - z) above
            d = np.where(r < Qn, Qn - z, np.where(r > Qp, Qp - z, q - r))
            ds = float((g * d).sum()) * grad_scale
            scale._accum(np.array([ds]))

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("cumsum backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            a._accum(np.flip(np.cumsum(np.flip(out.grad, axis=axis), axis=axis), axis=axis))

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("flip backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            a._accum(np.flip(out.grad, axis=axis))

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("logsumexp backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            g_e = out.grad if keepdims else np.expand_dims(out.grad, axis=axis)
            a._accum(sm * g_e)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("logaddexp backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        g = out.grad
        wa = np.exp(a.data - out_data)
        wb = np.exp(b.data - out_data)
        if a.requires_grad:
            a._accum(_unbroadcast(g * wa, a.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g * wb, b.shape))

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("l2_normalize backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if a.requires_grad:
            dot = (out.grad * y).sum(axis=axis, keepdims=True)
            dx = (out.grad - y * dot) / nrm
            a._accum(dx)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("rms_norm backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if x.requires_grad:
            gg = out.grad * gamma.data
            dx = inv * gg - (x.data * inv ** 3 / D) * (gg * x.data).sum(axis=-1, keepdims=True)
            x._accum(dx)
        if gamma.requires_grad:
            leading = tuple(range(out.grad.ndim - 1))
            gamma._accum((out.grad * xhat).sum(axis=leading))

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("instance_norm backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        g = out.grad
        if gamma.requires_grad:
            gamma._accum((g * xhat).sum(axis=(0, 2, 3)))
        if beta.requires_grad:
            beta._accum(g.sum(axis=(0, 2, 3)))
        if x.requires_grad:
            gx = (g * gamma.data.reshape(1, C, 1, 1)).reshape(N, C, M)
            inv_nc = inv  # shape (N, C, 1)
            xhat_nc = xhat.reshape(N, C, M)
            dx = inv_nc / M * (M * gx - gx.sum(axis=-1, keepdims=True) - xhat_nc * (gx * xhat_nc).sum(axis=-1, keepdims=True))
            x._accum(dx.reshape(N, C, H, W))

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("huber_loss backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if pred.requires_grad:
            dpred = np.where(quad, diff, delta * np.sign(diff)) / n * out.grad
            pred._accum(dpred)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("kl_div backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if log_p.requires_grad:
            log_p._accum(-q / n * out.grad)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("embedding backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if weight.requires_grad:
            dw = np.zeros_like(weight.data)
            np.add.at(dw, idx, out.grad)
            weight._accum(dw)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("conv2d_gen backward (grouped/dilated dW/db/dx)")  # TODO

    out._backward = _bw
    return out

    def _bw():
        g = out.grad.reshape(N, groups, cog, OH * OW)
        if weight.requires_grad:
            dW = np.einsum("ngop,ngcp->goc", g, cols_g).reshape(Cout, cig, KH, KW)
            weight._accum(dW)
        if has_bias and bias.requires_grad:
            bias._accum(out.grad.reshape(N, Cout, OH * OW).sum(axis=(0, 2)))
        if x.requires_grad:
            dcols_g = np.einsum("goc,ngop->ngcp", Wm, g)
            dcols = dcols_g.reshape(N, Cin * KH * KW, OH * OW)
            dxp = _col2im_dil(dcols, xp.shape, KH, KW, stride, dilation, OH, OW)
            dx = dxp[:, :, pad:-pad, pad:-pad] if pad > 0 else dxp
            x._accum(dx)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("conv_transpose2d backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        g_out = out.grad  # (N, Cout, OH, OW) - may be cropped
        # restore the full (un-cropped) gradient
        if pad > 0:
            g_full = np.zeros((N, Cout, OHf, OWf), dtype=np.float64)
            g_full[:, :, pad:OHf - pad, pad:OWf - pad] = g_out
        else:
            g_full = g_out
        # Gather patches: gcontrib[n,cout,i,j,ki,kj] = g_full[n,cout,i*stride+ki,j*stride+kj]
        gcontrib = np.zeros((N, Cout, H, W, KH, KW), dtype=np.float64)
        for ki in range(KH):
            for kj in range(KW):
                gcontrib[:, :, :, :, ki, kj] = g_full[:, :, ki:ki + stride * H:stride, kj:kj + stride * W:stride]
        if x.requires_grad:
            # dx[n,cin,i,j] = sum_{cout,ki,kj} Wm[cin,cout,ki,kj] * gcontrib[n,cout,i,j,ki,kj]
            dx = np.einsum("coKL,noijKL->ncij", Wm, gcontrib)
            x._accum(dx)
        if weight.requires_grad:
            # dW[cin,cout,ki,kj] = sum_{n,i,j} xd[n,cin,i,j] * gcontrib[n,cout,i,j,ki,kj]
            dW = np.einsum("ncij,noijKL->coKL", xd, gcontrib)
            weight._accum(dW)
        if has_bias and bias.requires_grad:
            bias._accum(g_out.sum(axis=(0, 2, 3)))

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("avgpool2d_s backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if x.requires_grad:
            dxp = np.zeros_like(xp)
            for oi in range(OH):
                for oj in range(OW):
                    dxp[:, :, oi * stride:oi * stride + k, oj * stride:oj * stride + k] += out.grad[:, :, oi:oi + 1, oj:oj + 1] / (k * k)
            if pad > 0:
                dx = dxp[:, :, pad:-pad, pad:-pad]
            else:
                dx = dxp
            x._accum(dx)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("maxpool2d_s backward")  # TODO

    out._backward = _bw
    return out

    def _bw():
        if x.requires_grad:
            dxp = np.zeros_like(xp)
            n_range = np.arange(N)[:, None, None, None]
            c_range = np.arange(C)[None, :, None, None]
            hi = np.arange(OH)[None, None, :, None] * stride + argi
            wj = np.arange(OW)[None, None, None, :] * stride + argj
            np.add.at(dxp, (n_range, c_range, hi, wj), out.grad)
            if pad > 0:
                dx = dxp[:, :, pad:-pad, pad:-pad]
            else:
                dx = dxp
            x._accum(dx)

    out._backward = _bw
    return out

    def _bw():
        raise NotImplementedError("fake_quant_lsq_plus backward (STE v-grad + beta-grad + LSQ scale grad)")  # TODO

    out._backward = _bw
    return out

    def _bw():
        g = out.grad
        ste_mask = (r > Qn) & (r < Qp)
        clip_mask = (r <= Qn) | (r >= Qp)
        if v.requires_grad:
            v._accum(g * ste_mask)
        if beta.requires_grad:
            beta._accum(np.array([float((g * clip_mask).sum())]))
        if scale.requires_grad:
            d = np.where(r < Qn, Qn, np.where(r > Qp, Qp, q - r))
            ds = float((g * d).sum()) * grad_scale
            scale._accum(np.array([ds]))

    out._backward = _bw
    return out

    out._backward = _bw
    return out


# ============================================================ LSQ fake-quant (DOMAIN; the crux)
def fake_quant(v, scale, n_levels_neg: int, n_levels_pos: int,

    out._backward = _bw
    return out


def batchnorm1d(x, gamma, beta, eps: float = 1e-5,
                running_mean=None, running_var=None, momentum: float = 0.1,
                training: bool = True) -> Tensor:
    """BatchNorm over N per channel C. x:(N,C); gamma,beta:(C,).

    TRAINING: normalize with the BATCH mean/var (POPULATION variance, divide by M=N); update
      running_mean/running_var IN PLACE if given (running var tracks the UNBIASED batch variance,
      i.e. var * M/(M-1)). The forward computes mu/var/inv/xhat; implements the backward.
    EVAL (training=False): normalize with the running stats; backward dx = grad*gamma*inv.
    """
    x, gamma, beta = _ensure(x), _ensure(gamma), _ensure(beta)
    N, C = x.data.shape
    M = N
    xd = x.data
    if training:
        mu = xd.mean(axis=0, keepdims=True)
        xc = xd - mu
        var = (xc * xc).mean(axis=0, keepdims=True)
        inv = 1.0 / np.sqrt(var + eps)
        xhat = xc * inv
        if running_mean is not None and running_var is not None:
            ubvar = var * (M / _bi_max(1, M - 1))
            running_mean[:] = (1 - momentum) * running_mean + momentum * mu.reshape(-1)
            running_var[:] = (1 - momentum) * running_var + momentum * ubvar.reshape(-1)
    else:
        mu = running_mean.reshape(1, C)
        var = running_var.reshape(1, C)
        inv = 1.0 / np.sqrt(var + eps)
        xhat = (xd - mu) * inv
    out = _mk(gamma.data * xhat + beta.data, (x, gamma, beta),
           
…[truncated 728 chars]

    def forward(self, x: Tensor) -> Tensor:
        # TODO: y = x @ weight (+ bias). Use ag.matmul / ag.add.
        raise NotImplementedError("Linear.forward")

    def forward(self, x: Tensor) -> Tensor:
        if self.bias is not None:
            return ag.add(ag.matmul(x, self.weight), self.bias)
        return ag.matmul(x, self.weight)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: out = x + fc2(gelu(fc1(x)))
        raise NotImplementedError("MLP.forward")

    def forward(self, x: Tensor) -> Tensor:
        return x + self.fc2(ag.gelu(self.fc1(x)))

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.layernorm(x, self.gamma, self.beta, self.eps)
        raise NotImplementedError("LayerNorm.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.layernorm(x, self.gamma, self.beta, self.eps)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.conv2d(x, self.weight, bias=self.bias, stride=self.stride, pad=self.padding)
        raise NotImplementedError("Conv2d.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.conv2d(x, self.weight, bias=self.bias, stride=self.stride, pad=self.padding)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.maxpool2d(x, self.k)
        raise NotImplementedError("MaxPool2d.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.maxpool2d(x, self.k)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.avgpool2d(x, self.k)
        raise NotImplementedError("AvgPool2d.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.avgpool2d(x, self.k)

    def quant(self, w: Tensor) -> Tensor:
        # TODO: return ag.fake_quant(w, self.scale, W_NEG, W_POS)
        raise NotImplementedError("WeightQuant.quant")

    def quant(self, w: Tensor) -> Tensor:
        return ag.fake_quant(w, self.scale, W_NEG, W_POS)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.fake_quant(x, self.scale, A_NEG, A_POS)
        raise NotImplementedError("ActQuant.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.fake_quant(x, self.scale, A_NEG, A_POS)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: ag.batchnorm2d(x, gamma, beta, eps, running_mean, running_var, momentum, training)
        raise NotImplementedError("BatchNorm2d.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.batchnorm2d(x, self.gamma, self.beta, self.eps, self.running_mean, self.running_var, self.momentum, self.training)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: ag.batchnorm1d(x, gamma, beta, eps, running_mean, running_var, momentum, training)
        raise NotImplementedError("BatchNorm1d.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.batchnorm1d(x, self.gamma, self.beta, self.eps, self.running_mean, self.running_var, self.momentum, self.training)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: wq = self.wq.quant(self.weight); return ag.conv2d(x, wq, bias=None, stride, pad=1)
        raise NotImplementedError("QConv.forward")

    def forward(self, x: Tensor) -> Tensor:
        wq = self.wq.quant(self.weight)
        return ag.conv2d(x, wq, bias=None, stride=self.stride, pad=1)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: h = aq(relu(bn(conv(x)))) -- in that order.
        raise NotImplementedError("ConvBlock.forward")

    def forward(self, x: Tensor) -> Tensor:
        h = self.conv(x)
        h = self.bn(h)
        h = ag.relu(h)
        h = self.aq(h)
        return h

    def forward(self, X) -> Tensor:
        # TODO: run the stem + blocks, global-average-pool, then the (2-bit) linear head.
        raise NotImplementedError("QuantCNN.forward")

    def forward(self, X) -> Tensor:
        x = Tensor(np.asarray(X, dtype=np.float64)) if not isinstance(X, Tensor) else X
        h = self.stem(x)
        for b in self.blocks:
            h = b(h)
        N, C, H, W = h.data.shape
        h = ag.avgpool2d(h, H)
        h = ag.reshape(h, (N, C))
        wq = self.head_wq.quant(self.head_w)
        logits = ag.add(ag.matmul(h, wq), self.head_b)
        return logits

    def forward(self, x: Tensor) -> Tensor:
        # TODO: wq = self.wq.quant(self.weight); return ag.add(ag.matmul(x, wq), self.bias)
        raise NotImplementedError("QLinear.forward")

    def forward(self, x: Tensor) -> Tensor:
        wq = self.wq.quant(self.weight)
        return ag.add(ag.matmul(x, wq), self.bias)

    def forward(self, X) -> Tensor:
        # TODO: x -> reshape (N, in_dim); for each layer L: h = L["aq"](relu(L["bn"](L["fc"](h))));
        # then return self.head(h).
        raise NotImplementedError("QuantMLP.forward")

    def forward(self, X) -> Tensor:
        x = X if isinstance(X, Tensor) else Tensor(np.asarray(X, dtype=np.float64))
        N = x.data.shape[0]
        h = ag.reshape(x, (N, self.in_dim))
        for L in self.layers:
            h = L["fc"](h)
            h = L["bn"](h)
            h = ag.relu(h)
            h = L["aq"](h)
        return self.head(h)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.silu(x)
        raise NotImplementedError("SiLU.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.silu(x)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.elu(x, self.alpha)
        raise NotImplementedError("ELU.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.elu(x, self.alpha)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.groupnorm(x, self.gamma, self.beta, self.num_groups, self.eps)
        raise NotImplementedError("GroupNorm.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.groupnorm(x, self.gamma, self.beta, self.num_groups, self.eps)

    def forward(self, x: Tensor, h: Tensor) -> Tensor:
        # TODO: return ag.tanh(x@Wx + h@Wh + b)
        raise NotImplementedError("RNNCell.forward")

    def forward(self, x: Tensor, h: Tensor) -> Tensor:
        return ag.tanh(ag.add(ag.add(ag.matmul(x, self.Wx), ag.matmul(h, self.Wh)), self.b))

    def forward(self, x: Tensor, h: Tensor) -> Tensor:
        # TODO: implement the GRU update with the named weights above.
        raise NotImplementedError("GRUCell.forward")

    def forward(self, x: Tensor, h: Tensor) -> Tensor:
        z = ag.sigmoid(ag.add(ag.add(ag.matmul(x, self.Wxz), ag.matmul(h, self.Whz)), self.bz))
        r = ag.sigmoid(ag.add(ag.add(ag.matmul(x, self.Wxr), ag.matmul(h, self.Whr)), self.br))
        n = ag.tanh(ag.add(ag.add(ag.matmul(x, self.Wxn), ag.matmul(ag.mul(r, h), self.Whn)), self.bn))
        h_new = ag.add(ag.mul(ag.sub(Tensor(np.ones_like(z.data)), z), n), ag.mul(z, h))
        return h_new

    def forward(self, x: Tensor, h: Tensor, c: Tensor):
        # TODO: implement the LSTM update; return (h2, c2).
        raise NotImplementedError("LSTMCell.forward")

    def forward(self, x: Tensor, h: Tensor, c: Tensor):
        i = ag.sigmoid(ag.add(ag.add(ag.matmul(x, self.Wxi), ag.matmul(h, self.Whi)), self.bi))
        f = ag.sigmoid(ag.add(ag.add(ag.matmul(x, self.Wxf), ag.matmul(h, self.Whf)), self.bf))
        g = ag.tanh(ag.add(ag.add(ag.matmul(x, self.Wxg), ag.matmul(h, self.Whg)), self.bg))
        o = ag.sigmoid(ag.add(ag.add(ag.matmul(x, self.Wxo), ag.matmul(h, self.Who)), self.bo))
        c2 = ag.add(ag.mul(f, c), ag.mul(i, g))
        h2 = ag.mul(o, ag.tanh(c2))
        return h2, c2

    def quant(self, w: Tensor) -> Tensor:
        # TODO: return ag.fake_quant_per_channel(w, self.scale, self.neg, self.pos, axis=0)
        raise NotImplementedError("PerChannelWeightQuant.quant")

    def quant(self, w: Tensor) -> Tensor:
        return ag.fake_quant_per_channel(w, self.scale, self.neg, self.pos, axis=0)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.fake_quant_affine(x, self.scale, self.zero_point, self.neg, self.pos)
        raise NotImplementedError("AffineActQuant.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.fake_quant_affine(x, self.scale, self.zero_point, self.neg, self.pos)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.rms_norm(x, self.gamma, self.eps)
        raise NotImplementedError("RMSNorm.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.rms_norm(x, self.gamma, self.eps)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.instance_norm(x, self.gamma, self.beta, self.eps)
        raise NotImplementedError("InstanceNorm2d.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.instance_norm(x, self.gamma, self.beta, self.eps)

    def forward(self, idx) -> Tensor:
        # TODO: return ag.embedding(self.weight, idx)
        raise NotImplementedError("Embedding.forward")

    def forward(self, idx) -> Tensor:
        return ag.embedding(self.weight, idx)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.conv_transpose2d(x, self.weight, bias=self.bias, stride=self.stride, pad=self.padding)
        raise NotImplementedError("ConvTranspose2d.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.conv_transpose2d(x, self.weight, bias=self.bias, stride=self.stride, pad=self.padding)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.conv2d_gen(x, self.weight, bias=self.bias, stride=self.stride,
        #                            pad=self.padding, dilation=self.dilation, groups=1)
        raise NotImplementedError("DilatedConv2d.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.conv2d_gen(x, self.weight, bias=self.bias, stride=self.stride, pad=self.padding, dilation=self.dilation, groups=1)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: return ag.fake_quant_lsq_plus(x, self.scale, self.beta, self.neg, self.pos)
        raise NotImplementedError("LSQPlusActQuant.forward")

    def forward(self, x: Tensor) -> Tensor:
        return ag.fake_quant_lsq_plus(x, self.scale, self.beta, self.neg, self.pos)

    def forward(self, x: Tensor) -> Tensor:
        # TODO: depthwise conv2d_gen(groups=Cin) on the LSQ-quantized dw_weight, then pointwise
        #       conv2d_gen(groups=1) on the LSQ-quantized pw_weight.
        raise NotImplementedError("QDWSepConv.forward")

    def forward(self, x: Tensor) -> Tensor:
        dw_q = self.dw_wq.quant(self.dw_weight)
        h = ag.conv2d_gen(x, dw_q, bias=None, stride=self.stride, pad=1, dilation=1, groups=self.Cin)
        pw_q = self.pw_wq.quant(self.pw_weight)
        h = ag.conv2d_gen(h, pw_q, bias=None, stride=1, pad=0, dilation=1, groups=1)
        return h

    def forward(self, x: Tensor) -> Tensor:
        # TODO: conv -> bn -> relu -> aq
        raise NotImplementedError("QSepBlock.forward")

    def forward(self, x: Tensor) -> Tensor:
        h = self.conv(x)
        h = self.bn(h)
        h = ag.relu(h)
        h = self.aq(h)
        return h

    def forward(self, X) -> Tensor:
        # TODO: stem -> blocks -> avgpool2d(full window) -> reshape (N,C) -> matmul(LSQ head) + bias
        raise NotImplementedError("QuantSepCNN.forward")

    def forward(self, X) -> Tensor:
        x = X if isinstance(X, Tensor) else Tensor(np.asarray(X, dtype=np.float64))
        h = self.stem(x)
        for b in self.blocks:
            h = b(h)
        N, C, H, W = h.data.shape
        h = ag.avgpool2d(h, H)
        h = ag.reshape(h, (N, C))
        wq = self.head_wq.quant(self.head_w)
        logits = ag.add(ag.matmul(h, wq), self.head_b)
        return logits

def normalize(X: np.ndarray) -> np.ndarray:
    """uint8 (N,3,S,S) -> float64, x = (X/255 - NORM_MEAN) / NORM_STD."""
    # TODO: cast to float, scale to [0,1], subtract NORM_MEAN, divide by NORM_STD.
    raise NotImplementedError("normalize")

def normalize(X: np.ndarray) -> np.ndarray:
    """uint8 (N,3,S,S) -> float64, x = (X/255 - NORM_MEAN) / NORM_STD."""
    return (X.astype(np.float64) / 255.0 - NORM_MEAN) / NORM_STD

    def step(self):
        # TODO: SGD(+momentum, +per-group weight decay) update. Iterate self._groups; for each
        # group use wd = self.weight_decay if override is None else override; keep a momentum
        # buffer per parameter (self._v is indexed in self.params order).
        raise NotImplementedError("SGD.step")

    def step(self):
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad + wd * p.data
                self._v[idx] = self.momentum * self._v[idx] + g
                p.data -= self.lr * self._v[idx]
                idx += 1

def iter_minibatches(Xn: np.ndarray, y: np.ndarray, batch_size: int,
                     rng: np.random.Generator, shuffle: bool = True):
    """Yield (xb, yb) minibatches over one epoch. ``Xn`` is already normalized float."""
    # TODO: optionally shuffle an index array, then yield contiguous batch_size slices of it.
    raise NotImplementedError("iter_minibatches")

def iter_minibatches(Xn: np.ndarray, y: np.ndarray, batch_size: int,
                     rng: np.random.Generator, shuffle: bool = True):
    """Yield (xb, yb) minibatches over one epoch. ``Xn`` is already normalized float."""
    N = len(Xn)
    idx = np.arange(N)
    if shuffle:
        rng.shuffle(idx)
    for start in range(0, N, batch_size):
        end = min(start + batch_size, N)
        batch_idx = idx[start:end]
        yield Xn[batch_idx], y[batch_idx]

    def step(self):
        # TODO: AdamW with bias correction + DECOUPLED, per-group weight decay (see docstring).
        raise NotImplementedError("AdamW.step")

    def step(self):
        self.t += 1
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad
                self._m[idx] = self.b1 * self._m[idx] + (1 - self.b1) * g
                self._v[idx] = self.b2 * self._v[idx] + (1 - self.b2) * g * g
                mhat = self._m[idx] / (1 - self.b1**self.t)
                vhat = self._v[idx] / (1 - self.b2**self.t)
                p.data -= self.lr * (mhat / (np.sqrt(vhat) + self.eps) + wd * p.data)
                idx += 1

    def step(self):
        # TODO: Adam with bias correction + COUPLED, per-group weight decay (see docstring).
        raise NotImplementedError("Adam.step")

    def step(self):
        self.t += 1
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad + wd * p.data
                self._m[idx] = self.b1 * self._m[idx] + (1 - self.b1) * g
                self._v[idx] = self.b2 * self._v[idx] + (1 - self.b2) * g * g
                mhat = self._m[idx] / (1 - self.b1**self.t)
                vhat = self._v[idx] / (1 - self.b2**self.t)
                p.data -= self.lr * mhat / (np.sqrt(vhat) + self.eps)
                idx += 1

    def step(self):
        # TODO: RMSprop update with per-group weight decay (see docstring).
        raise NotImplementedError("RMSprop.step")

    def step(self):
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad + wd * p.data
                self._s[idx] = self.alpha * self._s[idx] + (1 - self.alpha) * g * g
                p.data -= self.lr * g / (np.sqrt(self._s[idx]) + self.eps)
                idx += 1

    def step(self):
        # TODO: Adagrad update with per-group weight decay (see docstring).
        raise NotImplementedError("Adagrad.step")

    def step(self):
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad + wd * p.data
                self._s[idx] += g * g
                p.data -= self.lr * g / (np.sqrt(self._s[idx]) + self.eps)
                idx += 1

    def step(self):
        # TODO: Lion update with per-group (decoupled) weight decay (see docstring).
        raise NotImplementedError("Lion.step")

    def step(self):
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad
                update = np.sign(self.b1 * self._m[idx] + (1 - self.b1) * g)
                p.data -= self.lr * (update + wd * p.data)
                self._m[idx] = self.b2 * self._m[idx] + (1 - self.b2) * g
                idx += 1

def cosine_lr(step, total, base_lr, warmup=0, min_lr=0.0):
    """Linear warmup for ``warmup`` steps (lr = base_lr*(step+1)/warmup), then cosine decay from
    ``base_lr`` to ``min_lr`` over the remaining (total - warmup) steps. Return the LR at ``step``."""
    # TODO
    raise NotImplementedError("cosine_lr")

def cosine_lr(step, total, base_lr, warmup=0, min_lr=0.0):
    """Linear warmup for ``warmup`` steps (lr = base_lr*(step+1)/warmup), then cosine decay from
    ``base_lr`` to ``min_lr`` over the remaining (total - warmup) steps. Return the LR at ``step``."""
    if warmup > 0 and step < warmup:
        return base_lr * (step + 1) / warmup
    remaining = total - warmup
    t = min(step - warmup, remaining)
    progress = t / max(1, remaining)
    return min_lr + 0.5 * (base_lr - min_lr) * (1 + np.cos(np.pi * progress))

def warmup_lr(step, base_lr, warmup):
    """Linear warmup to ``base_lr`` over ``warmup`` steps (lr = base_lr*(step+1)/warmup), then
    constant ``base_lr``."""
    # TODO
    raise NotImplementedError("warmup_lr")

def warmup_lr(step, base_lr, warmup):
    """Linear warmup to ``base_lr`` over ``warmup`` steps (lr = base_lr*(step+1)/warmup), then
    constant ``base_lr``."""
    if step < warmup:
        return base_lr * (step + 1) / warmup
    return base_lr

def step_lr(step, base_lr, step_size, gamma=0.1):
    """Step decay: base_lr * gamma ** (step // step_size)."""
    # TODO
    raise NotImplementedError("step_lr")

def step_lr(step, base_lr, step_size, gamma=0.1):
    """Step decay: base_lr * gamma ** (step // step_size)."""
    return base_lr * (gamma ** (step // step_size))

def exponential_lr(step, base_lr, gamma=0.97):
    """Exponential decay: base_lr * gamma ** step."""
    # TODO
    raise NotImplementedError("exponential_lr")

def exponential_lr(step, base_lr, gamma=0.97):
    """Exponential decay: base_lr * gamma ** step."""
    return base_lr * (gamma ** step)

def clip_grad_norm(params: List[Tensor], max_norm: float) -> float:
    """GLOBAL L2 gradient-norm clip. Returns the PRE-CLIP total norm. If it exceeds ``max_norm``,
    scale every grad by ``max_norm / (total_norm + 1e-6)``."""
    # TODO: compute the global L2 norm over all param grads; scale in place if it exceeds max_norm.
    raise NotImplementedError("clip_grad_norm")

def clip_grad_norm(params: List[Tensor], max_norm: float) -> float:
    """GLOBAL L2 gradient-norm clip. Returns the PRE-CLIP total norm. If it exceeds ``max_norm``,
    scale every grad by ``max_norm / (total_norm + 1e-6)``."""
    total_norm = np.sqrt(sum(np.sum(p.grad**2) for p in params if p.grad is not None))
    if total_norm > max_norm:
        scale = max_norm / (total_norm + 1e-6)
        for p in params:
            if p.grad is not None:
                p.grad *= scale
    return float(total_norm)

    def step(self):
        # TODO: implement the Adamax update honouring per-group weight decay.
        raise NotImplementedError("Adamax.step")

    def step(self):
        self.t += 1
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad + wd * p.data
                self._m[idx] = self.b1 * self._m[idx] + (1 - self.b1) * g
                self._u[idx] = np.maximum(self.b2 * self._u[idx], np.abs(g))
                bias_corr = 1 - self.b1**self.t
                p.data -= (self.lr / bias_corr) * self._m[idx] / (self._u[idx] + self.eps)
                idx += 1

    def step(self):
        # TODO: implement the NAdam update honouring per-group weight decay.
        raise NotImplementedError("NAdam.step")

    def step(self):
        self.t += 1
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad + wd * p.data
                self._m[idx] = self.b1 * self._m[idx] + (1 - self.b1) * g
                self._v[idx] = self.b2 * self._v[idx] + (1 - self.b2) * g * g
                mhat = self.b1 * self._m[idx] / (1 - self.b1**(self.t+1)) + (1 - self.b1) * g / (1 - self.b1**self.t)
                vhat = self._v[idx] / (1 - self.b2**self.t)
                p.data -= self.lr * mhat / (np.sqrt(vhat) + self.eps)
                idx += 1

    def step(self):
        # TODO: implement the Adadelta update honouring per-group weight decay.
        raise NotImplementedError("Adadelta.step")

    def step(self):
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad + wd * p.data
                self._s[idx] = self.rho * self._s[idx] + (1 - self.rho) * g * g
                delta = np.sqrt(self._d[idx] + self.eps) / np.sqrt(self._s[idx] + self.eps) * g
                p.data -= self.lr * delta
                self._d[idx] = self.rho * self._d[idx] + (1 - self.rho) * delta * delta
                idx += 1

    def step(self):
        # TODO: implement the Nesterov-SGD update honouring per-group weight decay.
        raise NotImplementedError("NesterovSGD.step")

    def step(self):
        idx = 0
        for (group_params, wd_override) in self._groups:
            wd = self.weight_decay if wd_override is None else wd_override
            for p in group_params:
                if p.grad is None:
                    idx += 1
                    continue
                g = p.grad + wd * p.data
                self._v[idx] = self.momentum * self._v[idx] + g
                p.data -= self.lr * (g + self.momentum * self._v[idx])
                idx += 1

def polynomial_lr(step, total, base_lr, power=2.0, min_lr=0.0):
    """Polynomial decay: min_lr + (base_lr - min_lr) * (1 - step/total)^power (clamped at total)."""
    # TODO
    raise NotImplementedError("polynomial_lr")

def polynomial_lr(step, total, base_lr, power=2.0, min_lr=0.0):
    """Polynomial decay: min_lr + (base_lr - min_lr) * (1 - step/total)^power (clamped at total)."""
    t = min(step, total)
    return min_lr + (base_lr - min_lr) * (1 - t / total) ** power

def linear_lr(step, total, base_lr, min_lr=0.0):
    """Linear decay from base_lr to min_lr over ``total`` steps, then constant at min_lr."""
    # TODO
    raise NotImplementedError("linear_lr")

def linear_lr(step, total, base_lr, min_lr=0.0):
    """Linear decay from base_lr to min_lr over ``total`` steps, then constant at min_lr."""
    if step >= total:
        return min_lr
    return min_lr + (base_lr - min_lr) * (1 - step / total)

def multistep_lr(step, base_lr, milestones, gamma=0.1):
    """base_lr * gamma^(number of milestones <= step)."""
    # TODO
    raise NotImplementedError("multistep_lr")

def multistep_lr(step, base_lr, milestones, gamma=0.1):
    """base_lr * gamma^(number of milestones <= step)."""
    count = sum(1 for m in milestones if m <= step)
    return base_lr * (gamma ** count)

def cosine_restarts_lr(step, period, base_lr, min_lr=0.0):
    """Cosine annealing with warm restarts of fixed ``period`` (cosine within each period)."""
    # TODO
    raise NotImplementedError("cosine_restarts_lr")

def cosine_restarts_lr(step, period, base_lr, min_lr=0.0):
    """Cosine annealing with warm restarts of fixed ``period`` (cosine within each period)."""
    t = step % period
    progress = t / period
    return min_lr + 0.5 * (base_lr - min_lr) * (1 + np.cos(np.pi * progress))