qat-int2-cifar

    return grad.reshape(shape)


class Tensor:

    return grad.reshape(shape)


def _reduce_grad(grad, shape, axis, keepdims):
    """Broadcast a reduced gradient back to ``shape`` (re-inserting reduced axes if needed)."""
    if axis is None:
        return np.broadcast_to(grad, shape).copy()
    ax = axis if isinstance(axis, tuple) else (axis,)
    if not keepdims:
        grad = np.expand_dims(grad, ax)
    return np.broadcast_to(grad, shape).copy()


class Tensor:

        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")

        TODO: implement the topological walk + seeding. (The per-op ``_backward`` closures do the
        local accumulation; here you only orchestrate the order + seed ``self.grad``.)
        """
        topo: List["Tensor"] = []
        visited = set()

        def build(t):
            if id(t) in visited:
                return
            visited.add(id(t))
            for p in t._parents:
                if isinstance(p, Tensor):
                    build(p)
            topo.append(t)

        build(self)
        if grad is None:
            grad = np.ones_like(self.data)
        self._accum(_as_f64(grad))
        for t in reversed(topo):
            t._backward()

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

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

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

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

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

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

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

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

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

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

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

    def _bw():
        g = out.grad
        eq = a.data == b.data
        amask = np.where(a.data > b.data, 1.0, np.where(eq, 0.5, 0.0))
        bmask = 1.0 - amask
        if a.requires_grad:
            a._accum(_unbroadcast(g * amask, a.data.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g * bmask, b.data.shape))

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

    def _bw():
        g = out.grad
        eq = a.data == b.data
        amask = np.where(a.data < b.data, 1.0, np.where(eq, 0.5, 0.0))
        bmask = 1.0 - amask
        if a.requires_grad:
            a._accum(_unbroadcast(g * amask, a.data.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g * bmask, b.data.shape))

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

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

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

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

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

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

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

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

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

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

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

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

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

    def _bw():
        if a.requires_grad:
            a._accum(_reduce_grad(out.grad, a.data.shape, axis, keepdims))

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

    def _bw():
        if a.requires_grad:
            n = a.data.size / out.data.size
            a._accum(_reduce_grad(out.grad, a.data.shape, axis, keepdims) / n)

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

    def _bw():
        if a.requires_grad:
            xd = a.data
            m = xd.max(axis=axis, keepdims=True)
            mask = (xd == m)
            cnt = mask.sum(axis=axis, keepdims=True)
            g = out.grad if (keepdims or axis is None) else \
                np.expand_dims(out.grad, axis if isinstance(axis, tuple) else (axis,))
            a._accum(mask * (g / cnt))

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

    def _bw():
        if a.requires_grad:
            xd = a.data
            m = xd.min(axis=axis, keepdims=True)
            mask = (xd == m)
            cnt = mask.sum(axis=axis, keepdims=True)
            g = out.grad if (keepdims or axis is None) else \
                np.expand_dims(out.grad, axis if isinstance(axis, tuple) else (axis,))
            a._accum(mask * (g / cnt))

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

    def _bw():
        g = out.grad
        idx = 0
        for t in ts:
            sz = t.data.shape[axis]
            sl = [slice(None)] * g.ndim
            sl[axis] = slice(idx, idx + sz)
            if t.requires_grad:
                t._accum(g[tuple(sl)])
            idx += sz

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    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))

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

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

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

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

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

    def _bw():
        if logits.requires_grad:
            sm = np.exp(logp)
            grad = sm.copy()
            grad[np.arange(n), t] -= 1.0
            grad = grad / n * out.grad
            logits._accum(grad)

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

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

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

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

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

    def _bw():
        g = out.grad
        dout_mat = g.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:
            bias._accum(dout_mat.sum(axis=(0, 2)))
        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:H + pad, pad:W + pad] if pad > 0 else dxp
            x._accum(dx)

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

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

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

    def _bw():
        if x.requires_grad:
            g = out.grad
            m = xr.max(axis=(3, 5), keepdims=True)
            mask = (xr == m)
            cnt = mask.sum(axis=(3, 5), keepdims=True)
            gg = g.reshape(N, C, OH, 1, OW, 1)
            contrib = mask * (gg / cnt)
            full = np.zeros_like(x.data)
            full[:, :, :OH * k, :OW * k] = contrib.reshape(N, C, OH * k, OW * k)
            x._accum(full)

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

    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_
                sgx = gx.sum(axis=(0, 2, 3), keepdims=True)
                sgxx = (gx * xhat).sum(axis=(0, 2, 3), keepdims=True)
                dx = inv / M * (M * gx - sgx - xhat * sgxx)
            else:
                dx = g * g_ * inv
            x._accum(dx)

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

    def _bw():
        g = out.grad
        mask = (r >= Qn) & (r <= Qp)
        if v.requires_grad:
            v._accum(g * mask)
        if scale.requires_grad:
            ds_elem = np.where(mask, np.round(np.clip(r, Qn, Qp)) - r,
                               np.where(r < Qn, Qn, Qp))
            ds = float((g * ds_elem).sum()) * grad_scale
            scale._accum(np.full_like(scale.data, ds))

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

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

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

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

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

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

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

    def _bw():
        if a.requires_grad:
            n = a.data.size / out.data.size
            gb = _reduce_grad(out.grad, a.data.shape, axis, keepdims)
            a._accum(gb * (2.0 / n) * xc)

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

    def _bw():
        if a.requires_grad:
            n = a.data.size / out.data.size
            gb = _reduce_grad(out.grad, a.data.shape, axis, keepdims)
            a._accum(gb * xc / (n * sd))

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

    def _bw():
        if a.requires_grad:
            ga = np.zeros_like(a.data)
            grids = list(np.indices(out.data.shape))
            grids[axis] = np.broadcast_to(idx, out.data.shape)
            np.add.at(ga, tuple(grids), out.grad)
            a._accum(ga)

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

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

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

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

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

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

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

    def _bw():
        if a.requires_grad:
            th = np.tanh(sp)
            sig = 1.0 / (1.0 + np.exp(-x))
            a._accum(out.grad * (th + x * (1.0 - th ** 2) * sig))

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

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

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

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

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

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

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

    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)
            xh = xhat.reshape(N, G, M)
            sgx = gx.sum(axis=2, keepdims=True)
            sgxx = (gx * xh).sum(axis=2, keepdims=True)
            dx = inv / M * (M * gx - sgx - xh * sgxx)
            x._accum(dx.reshape(N, C, H, W))

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

    def _bw():
        g = out.grad
        mask = (r >= Qn) & (r <= Qp)
        if v.requires_grad:
            v._accum(g * mask)
        if scale.requires_grad:
            ds_elem = np.where(mask, np.round(np.clip(r, Qn, Qp)) - r,
                               np.where(r < Qn, Qn, Qp))
            axes = tuple(i for i in range(v.data.ndim) if i != axis)
            ds = (g * ds_elem).sum(axis=axes) * grad_scale
            scale._accum(ds.reshape(scale.data.shape))

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

    def _bw():
        g = out.grad
        mid = (r >= Qn) & (r <= Qp)
        if v.requires_grad:
            v._accum(g * mid)
        if scale.requires_grad:
            ds_elem = np.where(mid, q - r, np.where(r < Qn, Qn - z, Qp - z))
            ds = float((g * ds_elem).sum()) * grad_scale
            scale._accum(np.full_like(scale.data, ds))

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

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

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

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

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

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

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

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

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

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

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

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

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

    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 * g_).reshape(N, C, M)
            xh = xhat.reshape(N, C, M)
            sgx = gx.sum(axis=2, keepdims=True)
            sgxx = (gx * xh).sum(axis=2, keepdims=True)
            dx = inv / M * (M * gx - sgx - xh * sgxx)
            x._accum(dx.reshape(N, C, H, W))

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

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

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

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

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

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

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

    def _bw():
        g = out.grad.reshape(N, Cout, OH * OW)
        g_g = g.reshape(N, groups, cog, OH * OW)
        if weight.requires_grad:
            dW = np.einsum("ngop,ngcp->goc", g_g, cols_g).reshape(Cout, cig, KH, KW)
            weight._accum(dW)
        if has_bias and bias.requires_grad:
            bias._accum(g.sum(axis=(0, 2)))
        if x.requires_grad:
            dcols_g = np.einsum("goc,ngop->ngcp", Wm, g_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:H + pad, pad:W + pad] if pad > 0 else dxp
            x._accum(dx)

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

    def _bw():
        g = out.grad
        if has_bias and bias.requires_grad:
            bias._accum(g.sum(axis=(0, 2, 3)))
        if pad > 0:
            gfull = np.zeros((N, Cout, OHf, OWf), dtype=np.float64)
            gfull[:, :, pad:OHf - pad, pad:OWf - pad] = g
        else:
            gfull = g
        gcontrib = np.zeros((N, Cout, H, W, KH, KW), dtype=np.float64)
        for i in range(KH):
            for j in range(KW):
                gcontrib[:, :, :, :, i, j] = gfull[:, :, i:i + stride * H:stride, j:j + stride * W:stride]
        if x.requires_grad:
            dx = np.einsum("noijKL,coKL->ncij", gcontrib, Wm)
            x._accum(dx)
        if weight.requires_grad:
            dW = np.einsum("ncij,noijKL->coKL", xd, gcontrib)
            weight._accum(dW)

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

    def _bw():
        if x.requires_grad:
            g = out.grad
            dxp = np.zeros((N, C, Hp, Wp), dtype=np.float64)
            for oi in range(OH):
                for oj in range(OW):
                    dxp[:, :, oi * stride:oi * stride + k, oj * stride:oj * stride + k] += \
                        (g[:, :, oi, oj] / (k * k))[:, :, None, None]
            dx = dxp[:, :, pad:pad + H, pad:pad + W] if pad > 0 else dxp
            x._accum(dx)

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

    def _bw():
        if x.requires_grad:
            g = out.grad
            dxp = np.zeros((N, C, Hp, Wp), dtype=np.float64)
            nidx = np.arange(N)[:, None]
            cidx = np.arange(C)[None, :]
            for oi in range(OH):
                for oj in range(OW):
                    ii = oi * stride + argi[:, :, oi, oj]
                    jj = oj * stride + argj[:, :, oi, oj]
                    np.add.at(dxp, (nidx, cidx, ii, jj), g[:, :, oi, oj])
            dx = dxp[:, :, pad:pad + H, pad:pad + W] if pad > 0 else dxp
            x._accum(dx)

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

    def _bw():
        g = out.grad
        mid = (r >= Qn) & (r <= Qp)
        if v.requires_grad:
            v._accum(g * mid)
        if beta.requires_grad:
            db = float((g * (~mid)).sum())
            beta._accum(np.full_like(beta.data, db))
        if scale.requires_grad:
            ds_elem = np.where(mid, np.round(np.clip(r, Qn, Qp)) - r,
                               np.where(r < Qn, Qn, Qp))
            ds = float((g * ds_elem).sum()) * grad_scale
            scale._accum(np.full_like(scale.data, ds))

    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:
        y = ag.matmul(x, self.weight)
        if self.bias is not None:
            y = ag.add(y, self.bias)
        return y

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

    def forward(self, x: Tensor) -> Tensor:
        return ag.add(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)

    out._backward = _bw
    return out


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

    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 the batch axis N per feature C. x:(N,C); gamma,beta:(C,). Same semantics as
    batchnorm2d but reducing over axis 0 only (M=N)."""
    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
    g_ = gamma.data.reshape(1, C)
    out = _mk(g_ * xhat + beta.data.reshape(1, C), (x, gamma, beta),
              x.requires_grad or gamma.requires_grad or beta.requires_grad)

    def _bw():
        g = out.grad
        if gamma.requires_grad:
            gamma._accum((g * xhat).sum(axis=0))
        if beta.requires_grad:
            beta._accum(g.sum(axis=0))
…[truncated 471 chars]

    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:
        return self.aq(ag.relu(self.bn(self.conv(x))))

    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:
        h = X if isinstance(X, Tensor) else Tensor(X)
        h = self.stem(h)
        for b in self.blocks:
            h = b(h)
        N = h.data.shape[0]
        full_h = h.data.shape[2]
        pooled = ag.avgpool2d(h, full_h)
        pooled = ag.reshape(pooled, (N, self.feat))
        wq = self.head_wq.quant(self.head_w)
        return ag.add(ag.matmul(pooled, wq), self.head_b)

    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:
        h = X if isinstance(X, Tensor) else Tensor(X)
        N = h.data.shape[0]
        h = ag.reshape(h, (N, self.in_dim))
        for L in self.layers:
            h = L["aq"](ag.relu(L["bn"](L["fc"](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))
        one = Tensor(np.ones((1,)))
        return ag.add(ag.mul(ag.sub(one, z), n), ag.mul(z, h))

    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))
        gg = 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, gg))
        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:
        dwq = self.dw_wq.quant(self.dw_weight)
        h = ag.conv2d_gen(x, dwq, bias=None, stride=self.stride, pad=1, dilation=1, groups=self.Cin)
        pwq = self.pw_wq.quant(self.pw_weight)
        return ag.conv2d_gen(h, pwq, bias=None, stride=1, pad=0, dilation=1, groups=1)

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

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

    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:
        h = X if isinstance(X, Tensor) else Tensor(X)
        h = self.stem(h)
        for b in self.blocks:
            h = b(h)
        N = h.data.shape[0]
        full_h = h.data.shape[2]
        pooled = ag.reshape(ag.avgpool2d(h, full_h), (N, self.feat))
        wq = self.head_wq.quant(self.head_w)
        return ag.add(ag.matmul(pooled, wq), self.head_b)

    # TODO: cast to float, scale to [0,1], subtract NORM_MEAN, divide by NORM_STD.
    raise NotImplementedError("normalize")

    return (X.astype(np.float64) / 255.0 - NORM_MEAN) / NORM_STD

    # TODO: optionally shuffle an index array, then yield contiguous batch_size slices of it.
    raise NotImplementedError("iter_minibatches")

    n = Xn.shape[0]
    idx = np.arange(n)
    if shuffle:
        rng.shuffle(idx)
    for start in range(0, n, batch_size):
        sel = idx[start:start + batch_size]
        yield Xn[sel], y[sel]

    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):
        i = 0
        for group, override in self._groups:
            wd = self.weight_decay if override is None else override
            for p in group:
                if p.grad is not None:
                    g = p.grad + wd * p.data
                    self._v[i] = self.momentum * self._v[i] + g
                    p.data -= self.lr * self._v[i]
                i += 1

    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
        t = self.t
        i = 0
        for group, override in self._groups:
            wd = self.weight_decay if override is None else override
            for p in group:
                if p.grad is not None:
                    g = p.grad
                    self._m[i] = self.b1 * self._m[i] + (1 - self.b1) * g
                    self._v[i] = self.b2 * self._v[i] + (1 - self.b2) * (g * g)
                    mhat = self._m[i] / (1 - self.b1 ** t)
                    vhat = self._v[i] / (1 - self.b2 ** t)
                    p.data -= self.lr * (mhat / (np.sqrt(vhat) + self.eps) + wd * p.data)
                i += 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
        t = self.t
        i = 0
        for group, override in self._groups:
            wd = self.weight_decay if override is None else override
            for p in group:
                if p.grad is not None:
                    g = p.grad + wd * p.data
                    self._m[i] = self.b1 * self._m[i] + (1 - self.b1) * g
                    self._v[i] = self.b2 * self._v[i] + (1 - self.b2) * (g * g)
                    mhat = self._m[i] / (1 - self.b1 ** t)
                    vhat = self._v[i] / (1 - self.b2 ** t)
                    p.data -= self.lr * mhat / (np.sqrt(vhat) + self.eps)
                i += 1

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

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

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

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

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

    def step(self):
        i = 0
        for group, override in self._groups:
            wd = self.weight_decay if override is None else override
            for p in group:
                if p.grad is not None:
                    g = p.grad
                    upd = np.sign(self.b1 * self._m[i] + (1 - self.b1) * g)
                    p.data -= self.lr * (upd + wd * p.data)
                    self._m[i] = self.b2 * self._m[i] + (1 - self.b2) * g
                i += 1

    ``base_lr`` to ``min_lr`` over the remaining (total - warmup) steps. Return the LR at ``step``."""
    # TODO
    raise NotImplementedError("cosine_lr")

    ``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
    prog = min(1.0, max(0.0, (step - warmup) / max(1, total - warmup)))
    return min_lr + 0.5 * (base_lr - min_lr) * (1 + math.cos(math.pi * prog))

    constant ``base_lr``."""
    # TODO
    raise NotImplementedError("warmup_lr")

    constant ``base_lr``."""
    if step < warmup:
        return base_lr * (step + 1) / warmup
    return base_lr

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

    """Step decay: base_lr * gamma ** (step // step_size)."""
    return base_lr * gamma ** (step // step_size)

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

    """Exponential decay: base_lr * gamma ** step."""
    return base_lr * gamma ** step

    # TODO: compute the global L2 norm over all param grads; scale in place if it exceeds max_norm.
    raise NotImplementedError("clip_grad_norm")

    total = math.sqrt(sum(float((p.grad ** 2).sum()) for p in params if p.grad is not None))
    if total > max_norm:
        scale = max_norm / (total + 1e-6)
        for p in params:
            if p.grad is not None:
                p.grad = p.grad * scale
    return total

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

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

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

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

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

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

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

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

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

    """Polynomial decay: min_lr + (base_lr - min_lr) * (1 - step/total)^power (clamped at total)."""
    prog = min(1.0, max(0.0, step / total))
    return min_lr + (base_lr - min_lr) * ((1 - prog) ** power)

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

    """Linear decay from base_lr to min_lr over ``total`` steps, then constant at min_lr."""
    return base_lr + (min_lr - base_lr) * min(1.0, max(0.0, step / total))

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

    """base_lr * gamma^(number of milestones <= step)."""
    return base_lr * gamma ** sum(1 for m in milestones if step >= m)

    """Cosine annealing with warm restarts of fixed ``period`` (cosine within each period)."""
    # TODO
    raise NotImplementedError("cosine_restarts_lr")

    """Cosine annealing with warm restarts of fixed ``period`` (cosine within each period)."""
    t = step % period
    return min_lr + 0.5 * (base_lr - min_lr) * (1 + math.cos(math.pi * (t / period)))

def train(data_path, out_path, epochs=30, seed=0, n_classes=10):
    # TODO: load + normalize the data; build QuantCNN(n_classes, WIDTHS, seed); initialize the LSQ
    # step sizes; build the optimizer with weight decay only on the conv/head weights; run the QAT
    # loop (forward -> cross_entropy -> backward -> clip_grad_norm -> step) with an LR schedule;
    # keep your best epoch; BAKE the quantized weights; save the state_dict to out_path.
    raise NotImplementedError("implement the QAT training loop")

W_NEG, W_POS = 2, 1


def _init_scales(m):
    nps = m.named_params()
    for nm, p in nps.items():
        if nm.endswith("wq.scale"):
            wkey = nm[:-len("wq.scale")] + "weight"
            p.data[...] = max(2.0 * float(np.mean(np.abs(nps[wkey].data))) / np.sqrt(W_POS), 1e-2)
        elif nm.endswith("aq.scale"):
            p.data[...] = 1.0


def _copy_state(m):
    return {k: np.asarray(v).copy() for k, v in m.state_dict().items()}


def _evaluate(m, Xn, y, bs=256):
    m.set_training(False)
    correct = 0
    for i in range(0, len(y), bs):
        xb = Xn[i:i + bs]
        logits = m(xb).data
        correct += int((logits.argmax(1) == y[i:i + bs]).sum())
    m.set_training(True)
    return correct / len(y)


def _bake(m):
    """Replace each conv/head weight with its LSQ-quantized value on the 2-bit grid."""
    nps = m.named_params()
    for nm, p in nps.items():
        if nm.endswith(".weight") and (".conv" in nm or "stem" in nm or nm.startswith("head")
                                       or ".dw" in nm or ".pw" in nm) and p.data.ndim >= 2:
            skey = nm[:-len("weight")] + "wq.scale"
            if skey in nps:
                s = float(nps[skey].data.reshape(()))
                s = s if abs(s) > 1e-12 else 1e-12
                p.data[...] = np.round(np.clip(p.data / s, -W_NEG, W_POS)) * s


def train(data_path, out_path, epochs=60, seed=0, n_classes=10):
    rng = np.random.default_rng(1234)
    X, y = D.load_npz(data_path)
    Xn = D.normaliz
…[truncated 1627 chars]