fewlabel-ssl-fixmatch

        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: "Tensor"):
            if id(t) in visited:
                return
            visited.add(id(t))
            for p in t._parents:
                build(p)
            topo.append(t)

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

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

    out._backward = _bw
    return out


def sub(a, b) -> Tensor:

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

    out._backward = _bw
    return out


def sub(a, b) -> Tensor:

    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 * pf * (a.data ** (pf - 1.0)))

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

    def _bw():
        g = out.grad
        gt = (a.data > b.data).astype(np.float64)
        eq = (a.data == b.data).astype(np.float64) * 0.5
        sa = gt + eq
        if a.requires_grad:
            a._accum(_unbroadcast(g * sa, a.data.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g * (1.0 - sa), b.data.shape))

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

    def _bw():
        g = out.grad
        lt = (a.data < b.data).astype(np.float64)
        eq = (a.data == b.data).astype(np.float64) * 0.5
        sa = lt + eq
        if a.requires_grad:
            a._accum(_unbroadcast(g * sa, a.data.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(g * (1.0 - sa), b.data.shape))

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

    def _bw():
        g = out.grad
        cb = c.astype(bool)
        if a.requires_grad:
            a._accum(_unbroadcast(np.where(cb, g, 0.0), a.data.shape))
        if b.requires_grad:
            b._accum(_unbroadcast(np.where(cb, 0.0, g), 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)).astype(np.float64)
            a._accum(out.grad * mask)

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

    def _bw():
        if not a.requires_grad:
            return
        g = out.grad
        if axis is not None and not keepdims:
            g = np.expand_dims(g, axis)
        a._accum(np.broadcast_to(g, a.data.shape).copy())

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

    def _bw():
        if not a.requires_grad:
            return
        g = out.grad
        if axis is None:
            n = a.data.size
        elif isinstance(axis, tuple):
            n = int(np.prod([a.data.shape[ax] for ax in axis]))
        else:
            n = a.data.shape[axis]
        if axis is not None and not keepdims:
            g = np.expand_dims(g, axis)
        a._accum(np.broadcast_to(g / n, a.data.shape).copy())

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

    def _bw():
        if not a.requires_grad:
            return
        g = out.grad
        od = out.data
        if axis is not None and not keepdims:
            g = np.expand_dims(g, axis)
            od = np.expand_dims(od, axis)
        mask = (a.data == od).astype(np.float64)
        counts = mask.sum(axis=axis, keepdims=True)
        a._accum(mask * g / counts)

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

    def _bw():
        if not a.requires_grad:
            return
        g = out.grad
        od = out.data
        if axis is not None and not keepdims:
            g = np.expand_dims(g, axis)
            od = np.expand_dims(od, axis)
        mask = (a.data == od).astype(np.float64)
        counts = mask.sum(axis=axis, keepdims=True)
        a._accum(mask * g / counts)

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

    def _bw():
        if not a.requires_grad:
            return
        if axes is None:
            a._accum(np.transpose(out.grad))
        else:
            inv = np.argsort(axes)
            a._accum(np.transpose(out.grad, tuple(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 not a.requires_grad:
            return
        gz = np.zeros_like(a.data)
        np.add.at(gz, idx, out.grad)
        a._accum(gz)

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

    def _bw():
        g = out.grad
        offsets = np.cumsum([0] + [t.data.shape[axis] for t in ts])
        for i, t in enumerate(ts):
            if t.requires_grad:
                sl = [slice(None)] * g.ndim
                sl[axis] = slice(offsets[i], offsets[i + 1])
                t._accum(g[tuple(sl)])

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    def _bw():
        if a.requires_grad:
            grad_local = np.where(pos, 1.0, al * np.exp(np.minimum(x, 0.0)))
            a._accum(out.grad * grad_local)

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

    def _bw():
        if not a.requires_grad:
            return
        sm = e / s
        g = out.grad
        if not keepdims:
            g = np.expand_dims(g, axis)
        a._accum(g * sm)

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

    def _bw():
        if not a.requires_grad:
            return
        if axis is None:
            N = x.size
        elif isinstance(axis, tuple):
            N = int(np.prod([x.shape[ax] for ax in axis]))
        else:
            N = x.shape[axis]
        g = out.grad
        if axis is not None and not keepdims:
            g = np.expand_dims(g, axis)
        a._accum(g * (2.0 / N) * xc)

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

    def _bw():
        if not a.requires_grad:
            return
        if axis is None:
            N = x.size
        elif isinstance(axis, tuple):
            N = int(np.prod([x.shape[ax] for ax in axis]))
        else:
            N = x.shape[axis]
        g = out.grad
        if axis is not None and not keepdims:
            g = np.expand_dims(g, axis)
        a._accum(g * xc / (N * sd))

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

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

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

    def _bw():
        g = out.grad
        if a.requires_grad:
            a._accum((g @ b.data.reshape(-1)).reshape(a.data.shape))
        if b.requires_grad:
            b._accum((a.data.reshape(-1) @ g).reshape(b.data.shape))

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

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

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

    def _bw():
        if a.requires_grad:
            g = out.grad
            a._accum(s * (g - (g * s).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(ls)
            a._accum(g - sm * g.sum(axis=axis, keepdims=True))

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

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

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

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

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

    def _bw():
        g = out.grad
        D = x.shape[-1]
        if gamma.requires_grad:
            axes = tuple(range(g.ndim - 1))
            gamma._accum((g * xhat).sum(axis=axes))
        if beta.requires_grad:
            axes = tuple(range(g.ndim - 1))
            beta._accum(g.sum(axis=axes))
        if a.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))
            a._accum(dx)

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

    def _bw():
        g = out.grad
        N = x.shape[0]
        if gamma.requires_grad:
            gamma._accum((g * xhat).sum(axis=0))
        if beta.requires_grad:
            beta._accum(g.sum(axis=0))
        if a.requires_grad:
            gx = g * gamma.data
            dx = (inv / N) * (N * gx - gx.sum(axis=0, keepdims=True)
                              - xhat * (gx * xhat).sum(axis=0, keepdims=True))
            a._accum(dx)

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

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

        def _bw_id():
            raise NotImplementedError("dropout (eval) backward")  # TODO

        out._backward = _bw_id
        return out
    raise NotImplementedError("dropout (training) -- optional")

        def _bw_id():
            if a.requires_grad:
                a._accum(out.grad)

        out._backward = _bw_id
        return out
    if rng is None:
        rng = np.random.default_rng()
    mask = (rng.random(a.data.shape) >= p).astype(np.float64) / (1.0 - p)
    out = _mk(a.data * mask, (a,), a.requires_grad)

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

    out._backward = _bw
    return out

    TODO: implement the FORWARD and the backward. (Unlike the other ops, the forward is NOT
    pre-filled for you -- build the output, then differentiate it.)
    """
    raise NotImplementedError("conv2d")

    TODO: implement the FORWARD and the backward. (Unlike the other ops, the forward is NOT
    pre-filled for you -- build the output, then differentiate it.)
    """
    x = _ensure(x)
    weight = _ensure(weight)
    bias_t = _ensure(bias) if bias is not None else None
    xd = x.data
    wd = weight.data
    N, Cin, H, W = xd.shape
    Cout, _, KH, KW = wd.shape
    s, p = stride, padding
    OH = (H + 2 * p - KH) // s + 1
    OW = (W + 2 * p - KW) // s + 1
    xp = _pad2d(xd, p)
    # im2col -> (N, Cin*KH*KW, OH*OW)
    cols = np.empty((N, Cin * KH * KW, OH * OW), dtype=np.float64)
    for i in range(KH):
        for j in range(KW):
            patch = xp[:, :, i:i + s * OH:s, j:j + s * OW:s]  # (N,Cin,OH,OW)
            cols[:, (i * KW + j)::KH * KW, :] = 0.0  # placeholder, overwritten below
    # build cols properly ordered as (Cin, KH, KW)
    cols = np.empty((N, Cin, KH, KW, OH, OW), dtype=np.float64)
    for i in range(KH):
        for j in range(KW):
            cols[:, :, i, j, :, :] = xp[:, :, i:i + s * OH:s, j:j + s * OW:s]
    cols2 = cols.reshape(N, Cin * KH * KW, OH * OW)
    wmat = wd.reshape(Cout, Cin * KH * KW)
    out_data = np.einsum('oc,ncp->nop', wmat, cols2).reshape(N, Cout, OH, OW)
    if bias_t is not None:
        out_data = out_data + bias_t.data.reshape(1, Cout, 1, 1)
    parents = (x, weight) if bias_t is None else (x, weight, bias_t)
    rg = x.requires_grad or weight.requires_grad or (bias_t is not None and bias_t.requires_grad)
    out = _mk
…[truncated 904 chars]

    xp = _pad2d(xd, p)
    # im2col -> (N, Cin*KH*KW, OH*OW)
    cols = np.empty((N, Cin * KH * KW, OH * OW), dtype=np.float64)
    for i in range(KH):
        for j in range(KW):
            patch = xp[:, :, i:i + s * OH:s, j:j + s * OW:s]  # (N,Cin,OH,OW)
            cols[:, (i * KW + j)::KH * KW, :] = 0.0  # placeholder, overwritten below
    # build cols properly ordered as (Cin, KH, KW)
    cols = np.empty((N, Cin, KH, KW, OH, OW), dtype=np.float64)

    xp = _pad2d(xd, p)
    # im2col -> (N, Cin, KH, KW, OH, OW)
    cols = np.empty((N, Cin, KH, KW, OH, OW), dtype=np.float64)

    TODO: implement the FORWARD and the backward.
    """
    raise NotImplementedError("maxpool2d")

    TODO: implement the FORWARD and the backward.
    """
    x = _ensure(x)
    xd = x.data
    N, C, H, W = xd.shape
    k, s = kernel, stride
    OH = (H - k) // s + 1
    OW = (W - k) // s + 1
    win = np.empty((N, C, k, k, OH, OW), dtype=np.float64)
    for i in range(k):
        for j in range(k):
            win[:, :, i, j, :, :] = xd[:, :, i:i + s * OH:s, j:j + s * OW:s]
    win_r = win.reshape(N, C, k * k, OH, OW)
    out_data = win_r.max(axis=2)
    out = _mk(out_data, (x,), x.requires_grad)

    def _bw():
        if not x.requires_grad:
            return
        g = out.grad
        mx = out_data[:, :, None, :, :]
        mask = (win_r == mx).astype(np.float64)
        counts = mask.sum(axis=2, keepdims=True)
        contrib = mask / counts * g[:, :, None, :, :]  # (N,C,k*k,OH,OW)
        contrib = contrib.reshape(N, C, k, k, OH, OW)
        dx = np.zeros_like(xd)
        for i in range(k):
            for j in range(k):
                dx[:, :, i:i + s * OH:s, j:j + s * OW:s] += contrib[:, :, i, j, :, :]
        x._accum(dx)

    out._backward = _bw
    return out

    TODO: implement the FORWARD and the backward.
    """
    raise NotImplementedError("avgpool2d")

    TODO: implement the FORWARD and the backward.
    """
    x = _ensure(x)
    xd = x.data
    N, C, H, W = xd.shape
    k, s = kernel, stride
    OH = (H - k) // s + 1
    OW = (W - k) // s + 1
    acc = np.zeros((N, C, OH, OW), dtype=np.float64)
    for i in range(k):
        for j in range(k):
            acc += xd[:, :, i:i + s * OH:s, j:j + s * OW:s]
    out_data = acc / (k * k)
    out = _mk(out_data, (x,), x.requires_grad)

    def _bw():
        if not x.requires_grad:
            return
        g = out.grad / (k * k)
        dx = np.zeros_like(xd)
        for i in range(k):
            for j in range(k):
                dx[:, :, i:i + s * OH:s, j:j + s * OW:s] += g
        x._accum(dx)

    out._backward = _bw
    return out

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

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

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

        return ag.embedding(self.weight, idx)

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

        return ag.layernorm(x, self.gamma, self.beta, self.eps)

        # TODO: ms = mean(x*x, -1, keepdims); inv = (ms+eps)**-0.5; return x*inv*gamma  (autograd ops)
        raise NotImplementedError("RMSNorm.forward")

        ms = ag.mean(ag.mul(x, x), axis=-1, keepdims=True)
        inv = ag.power(ag.add(ms, self.eps), -0.5)
        return ag.mul(ag.mul(x, inv), self.gamma)

        # TODO: g = ag.silu(self.w_gate(x)); u = self.w_up(x); return self.w_down(ag.mul(g, u))
        raise NotImplementedError("SwiGLU.forward")

        g = ag.silu(self.w_gate(x))
        u = self.w_up(x)
        return self.w_down(ag.mul(g, u))

        # TODO: return ag.batchnorm(x, self.gamma, self.beta, self.eps)
        raise NotImplementedError("BatchNorm1d.forward")

        return ag.batchnorm(x, self.gamma, self.beta, self.eps)

        # TODO: return ag.dropout(x, self.p, training, rng)
        raise NotImplementedError("Dropout.forward")

        return ag.dropout(x, self.p, training, rng)

        # TODO: return ag.conv2d(x, self.weight, self.bias, self.stride, self.padding)
        raise NotImplementedError("Conv2d.forward")

        return ag.conv2d(x, self.weight, self.bias, self.stride, self.padding)

        # TODO: return ag.maxpool2d(x, self.kernel, self.stride)
        raise NotImplementedError("MaxPool2d.forward")

        return ag.maxpool2d(x, self.kernel, self.stride)

        # TODO: return ag.avgpool2d(x, self.kernel, self.stride)
        raise NotImplementedError("AvgPool2d.forward")

        return ag.avgpool2d(x, self.kernel, self.stride)

        # TODO: implement causal multi-head self-attention as described in the docstring.
        raise NotImplementedError("MultiHeadSelfAttention.forward")

        B, Tt, _ = x.shape
        H, dh = self.n_head, self.d_head
        qkv = self.qkv(x)  # (B,T,3*d_model)
        qkv = ag.reshape(qkv, (B, Tt, 3, H, dh))
        qkv = ag.transpose(qkv, (2, 0, 3, 1, 4))  # (3,B,H,T,dh)
        q = qkv[0]
        k = qkv[1]
        v = qkv[2]
        kt = ag.transpose(k, (0, 1, 3, 2))  # (B,H,dh,T)
        scores = ag.mul(ag.matmul(q, kt), 1.0 / np.sqrt(dh))  # (B,H,T,T)
        mask = np.triu(np.ones((Tt, Tt), dtype=np.float64), k=1) * (-1e9)
        scores = ag.add(scores, Tensor(mask.reshape(1, 1, Tt, Tt)))
        att = ag.softmax(scores, axis=-1)
        o = ag.matmul(att, v)  # (B,H,T,dh)
        o = ag.transpose(o, (0, 2, 1, 3))  # (B,T,H,dh)
        o = ag.reshape(o, (B, Tt, self.d_model))
        return self.proj(o)

        # TODO: fc2(gelu(fc1(x)))
        raise NotImplementedError("MLP.forward")

        return self.fc2(ag.gelu(self.fc1(x)))

        # TODO: pre-norm residual wiring (see docstring).
        raise NotImplementedError("TransformerBlock.forward")

        x = ag.add(x, self.attn(self.ln1(x)))
        x = ag.add(x, self.mlp(self.ln2(x)))
        return x

        # TODO: tanh(x@W_ih + b_ih + h@W_hh + b_hh)
        raise NotImplementedError("RNNCell.forward")

        pre = ag.add(ag.add(ag.matmul(x, self.W_ih), self.b_ih),
                     ag.add(ag.matmul(h, self.W_hh), self.b_hh))
        return ag.tanh(pre)

        # TODO: implement the LSTM cell update; return (h_new, c_new).
        raise NotImplementedError("LSTMCell.forward")

        h, c = state
        H = self.hidden
        z = ag.add(ag.add(ag.matmul(x, self.W_ih), self.b_ih),
                   ag.add(ag.matmul(h, self.W_hh), self.b_hh))
        i = ag.sigmoid(z[:, 0 * H:1 * H])
        f = ag.sigmoid(z[:, 1 * H:2 * H])
        g = ag.tanh(z[:, 2 * H:3 * H])
        o = ag.sigmoid(z[:, 3 * H:4 * H])
        c_new = ag.add(ag.mul(f, c), ag.mul(i, g))
        h_new = ag.mul(o, ag.tanh(c_new))
        return h_new, c_new

        # TODO: implement the GRU cell update; return h_new.
        raise NotImplementedError("GRUCell.forward")

        H = self.hidden
        gi = ag.add(ag.matmul(x, self.W_ih), self.b_ih)
        gh = ag.add(ag.matmul(h, self.W_hh), self.b_hh)
        r = ag.sigmoid(ag.add(gi[:, 0:H], gh[:, 0:H]))
        z = ag.sigmoid(ag.add(gi[:, H:2 * H], gh[:, H:2 * H]))
        n = ag.tanh(ag.add(gi[:, 2 * H:3 * H], ag.mul(r, gh[:, 2 * H:3 * H])))
        one_minus_z = ag.sub(1.0, z)
        return ag.add(ag.mul(one_minus_z, n), ag.mul(z, h))

        # TODO: relu(layernorm(linear(x))) -- compose self.fc, self.ln and ag.relu.
        raise NotImplementedError("HiddenBlock.forward")

        return ag.relu(self.ln(self.fc(x)))

        # TODO: h = Tensor(X) if X is a raw array; for blk in self.layers: h = blk(h);
        #       return self.head(h)
        raise NotImplementedError("MLPClassifier.forward")

        h = X if isinstance(X, Tensor) else Tensor(X)
        for blk in self.layers:
            h = blk(h)
        return self.head(h)

        # TODO: implement the conv->relu->maxpool x2 -> flatten -> fc pipeline (see docstring).
        #       Wrap a raw ndarray x in a Tensor first; reshape with the batch size N as dim 0.
        raise NotImplementedError("CNNClassifier.forward")

        h = x if isinstance(x, Tensor) else Tensor(x)
        N = h.shape[0]
        h = ag.maxpool2d(ag.relu(self.conv1(h)), 2, 2)
        h = ag.maxpool2d(ag.relu(self.conv2(h)), 2, 2)
        h = ag.reshape(h, (N, self.feat))
        return self.fc(h)

        # TODO: embed idx, loop the cell over time from a zero state, classify the final hidden.
        raise NotImplementedError("SeqClassifier.forward")

        idx = np.asarray(idx).astype(np.int64)
        N, T = idx.shape
        emb = self.embed(idx)  # (N,T,d_model)
        h = Tensor(np.zeros((N, self.hidden)))
        if self.cell_kind == "lstm":
            c = Tensor(np.zeros((N, self.hidden)))
            for t in range(T):
                h, c = self.cell(emb[:, t, :], (h, c))
        else:
            for t in range(T):
                h = self.cell(emb[:, t, :], h)
        return self.fc(h)

    def step(self):
        # TODO: SGD(+momentum, +coupled weight decay, +optional Nesterov) update.
        raise NotImplementedError("SGD.step")

    def step(self):
        for i, p in enumerate(self.params):
            if p.grad is None:
                continue
            g = p.grad + self.weight_decay * p.data
            v = self.momentum * self._v[i] + g
            self._v[i] = v
            upd = (g + self.momentum * v) if self.nesterov else v
            p.data -= self.lr * upd

        # TODO: Adagrad update (accumulate squared grads).
        raise NotImplementedError("Adagrad.step")

        for i, p in enumerate(self.params):
            if p.grad is None:
                continue
            g = p.grad + self.weight_decay * p.data
            self._s[i] += g * g
            p.data -= self.lr * g / (np.sqrt(self._s[i]) + self.eps)

        # TODO: RMSprop update (EMA of squared grads).
        raise NotImplementedError("RMSprop.step")

        for i, p in enumerate(self.params):
            if p.grad is None:
                continue
            g = p.grad + self.weight_decay * 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)

        # TODO: Adam update with bias correction + COUPLED (L2) weight decay.
        raise NotImplementedError("Adam.step")

        self.t += 1
        for i, p in enumerate(self.params):
            if p.grad is None:
                continue
            g = p.grad + self.weight_decay * 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 ** self.t)
            vhat = self._v[i] / (1 - self.b2 ** self.t)
            p.data -= self.lr * mhat / (np.sqrt(vhat) + self.eps)

        # TODO: AdamW update with bias correction + DECOUPLED weight decay.
        raise NotImplementedError("AdamW.step")

        self.t += 1
        for i, p in enumerate(self.params):
            if p.grad is None:
                continue
            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 ** self.t)
            vhat = self._v[i] / (1 - self.b2 ** self.t)
            p.data -= self.lr * (mhat / (np.sqrt(vhat) + self.eps) + self.weight_decay * p.data)

        # TODO: Lion update (sign of interpolated momentum; momentum updated AFTER the step).
        raise NotImplementedError("Lion.step")

        for i, p in enumerate(self.params):
            if p.grad is None:
                continue
            g = p.grad
            upd = np.sign(self.b1 * self._m[i] + (1 - self.b1) * g)
            p.data -= self.lr * (upd + self.weight_decay * p.data)
            self._m[i] = self.b2 * self._m[i] + (1 - self.b2) * g

    # 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:
        sc = max_norm / (total + 1e-6)
        for p in params:
            if p.grad is not None:
                p.grad = p.grad * sc
    return total

    # TODO
    raise NotImplementedError("cosine_lr")

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

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

    """Linear warmup to ``base_lr`` over ``warmup`` steps, then constant ``base_lr``."""
    # TODO
    raise NotImplementedError("warmup_lr")

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

def load_npz(path: str) -> Tuple[np.ndarray, Optional[np.ndarray]]:
    # TODO: np.load the file; return (X as float64, y as int64 reshaped to 1-D or None if absent).
    raise NotImplementedError("load_npz")


def n_classes(y: np.ndarray) -> int:
    # TODO: return int(max(y)) + 1
    raise NotImplementedError("n_classes")


def epoch_batches(n: int, batch_size: int, rng: np.random.Generator):
    # TODO: yield successive slices of a single rng.permutation(n) of size batch_size.
    raise NotImplementedError("epoch_batches")


def get_batch(X: np.ndarray, y: Optional[np.ndarray], idx: np.ndarray):
    # TODO: gather X[idx] (float64) and y[idx] (int64, or None if y is None); return (Xb, yb, idx).
    raise NotImplementedError("get_batch")

def load_npz(path: str) -> Tuple[np.ndarray, Optional[np.ndarray]]:
    d = np.load(path)
    X = np.asarray(d["X"], dtype=np.float64)
    y = None
    if "y" in d.files:
        y = np.asarray(d["y"], dtype=np.int64).reshape(-1)
    return X, y


def n_classes(y: np.ndarray) -> int:
    return int(np.max(y)) + 1


def epoch_batches(n: int, batch_size: int, rng: np.random.Generator):
    perm = rng.permutation(n)
    for i in range(0, n, batch_size):
        yield perm[i:i + batch_size]


def get_batch(X: np.ndarray, y: Optional[np.ndarray], idx: np.ndarray):
    idx = np.asarray(idx).astype(np.int64)
    Xb = np.asarray(X[idx], dtype=np.float64)
    yb = None if y is None else np.asarray(y[idx], dtype=np.int64)
    return Xb, yb, idx

    # TODO: build per-class templates, pick labels, return template[y] + noise and the labels.
    raise NotImplementedError("make_image_dataset")

    rng = np.random.default_rng(seed)
    templates = rng.standard_normal((n_classes, 1, img, img))
    y = rng.integers(0, n_classes, size=n)
    X = templates[y] + rng.standard_normal((n, 1, img, img)) * 0.3
    return np.asarray(X, dtype=np.float64), np.asarray(y, dtype=np.int64)

    # TODO: sample (n, T) token ids in [0, vocab); label = (row sum) % n_classes.
    raise NotImplementedError("make_seq_dataset")

    rng = np.random.default_rng(seed)
    X = rng.integers(0, vocab, size=(n, T))
    y = (X.sum(axis=1) % n_classes).astype(np.int64)
    return np.asarray(X, dtype=np.int64), y

    # TODO: return X + rng.normal(0, sigma, X.shape)  (X as float64)
    raise NotImplementedError("weak_augment")

    X = np.asarray(X, dtype=np.float64)
    return X + rng.normal(0.0, sigma, X.shape)

    # TODO: add jitter, then for each row zero round(mask_frac*dim) randomly-chosen coordinates.
    raise NotImplementedError("strong_augment")

    X = np.asarray(X, dtype=np.float64)
    out = X + rng.normal(0.0, sigma, X.shape)
    n, dim = out.shape
    k = int(round(mask_frac * dim))
    if k > 0:
        for r in range(n):
            cols = rng.choice(dim, size=k, replace=False)
            out[r, cols] = 0.0
    return out

    # TODO: copy X; for each row pick a random start and zero [start:start+length].
    raise NotImplementedError("cutout")

    out = np.array(X, dtype=np.float64)
    n, dim = out.shape
    for r in range(n):
        start = int(rng.integers(0, dim))
        end = min(start + length, dim)
        out[r, start:end] = 0.0
    return out

    # TODO: implement the convex combination described above; return (X_mix, y_mix, lam, perm).
    raise NotImplementedError("mixup")

    X = np.asarray(X, dtype=np.float64)
    onehot = np.asarray(onehot, dtype=np.float64)
    lam = float(rng.beta(alpha, alpha))
    lam = max(lam, 1.0 - lam)
    perm = rng.permutation(X.shape[0])
    X_mix = lam * X + (1.0 - lam) * X[perm]
    y_mix = lam * onehot + (1.0 - lam) * onehot[perm]
    return X_mix, y_mix, lam, perm

    # TODO: stack to (K, M, ...), swap the leading block axis with K over full blocks, return list.
    raise NotImplementedError("interleave")

    arrs = [np.asarray(a) for a in arrays]
    K = len(arrs)
    M = arrs[0].shape[0]
    rest = arrs[0].shape[1:]
    nb = M // batch_size
    s = np.stack(arrs, 0).reshape(K, nb, batch_size, *rest)
    s = np.swapaxes(s, 0, 1).reshape(K, M, *rest)
    return [s[i] for i in range(K)]

    # TODO: invert the interleave permutation.
    raise NotImplementedError("de_interleave")

    arrs = [np.asarray(a) for a in arrays]
    K = len(arrs)
    M = arrs[0].shape[0]
    rest = arrs[0].shape[1:]
    nb = M // batch_size
    s = np.stack(arrs, 0).reshape(nb, K, batch_size, *rest)
    s = np.swapaxes(s, 0, 1).reshape(K, M, *rest)
    return [s[i] for i in range(K)]

    # TODO: compute hard/conf/mask/onehot as above.
    raise NotImplementedError("pseudo_label")

    weak_probs = np.asarray(weak_probs, dtype=np.float64)
    n, c = weak_probs.shape
    hard = weak_probs.argmax(axis=1)
    conf = weak_probs.max(axis=1)
    mask = (conf >= tau).astype(np.float64)
    onehot = np.zeros((n, c), dtype=np.float64)
    onehot[np.arange(n), hard] = 1.0
    return onehot, mask, hard

    # TODO: return (probs.max(axis=1) >= tau).astype(float)
    raise NotImplementedError("confidence_mask")

    probs = np.asarray(probs, dtype=np.float64)
    return (probs.max(axis=1) >= tau).astype(np.float64)

    # TODO: pw = probs ** (1/temperature); return pw / pw.sum(axis=1, keepdims=True)
    raise NotImplementedError("sharpen")

    probs = np.asarray(probs, dtype=np.float64)
    pw = probs ** (1.0 / temperature)
    return pw / pw.sum(axis=1, keepdims=True)

        # TODO: update self.value with the EMA recursion and return it.
        raise NotImplementedError("EMAScalar.update")

        self.value = self.decay * self.value + (1.0 - self.decay) * float(new)
        return self.value

        # TODO: update self.running, rescale by prior/(running+1e-6), renormalize rows, return.
        raise NotImplementedError("DistributionAligner.align")

        weak_probs = np.asarray(weak_probs, dtype=np.float64)
        self.running = self.momentum * self.running + (1.0 - self.momentum) * weak_probs.mean(axis=0)
        aligned = weak_probs * (self.prior / (self.running + 1e-6))
        aligned = aligned / aligned.sum(axis=1, keepdims=True)
        return aligned

        # TODO: compute sigma (per-class confident-argmax counts), set self.sigma, return beta*tau
        # with denom = max(sigma.max(), C - sigma.sum()) (guard against 0).
        raise NotImplementedError("FlexMatchThresholder.update")

    def flex_mask(self, weak_probs) -> np.ndarray:
        # TODO: thr = self.update(weak_probs); return (max prob per row >= thr[argmax]) as 0/1 float.
        raise NotImplementedError("FlexMatchThresholder.flex_mask")

        weak_probs = np.asarray(weak_probs, dtype=np.float64)
        hard = weak_probs.argmax(axis=1)
        conf = weak_probs.max(axis=1)
        adm = conf >= self.tau
        sigma = np.array([float(np.sum(adm & (hard == c))) for c in range(self.C)],
                         dtype=np.float64)
        self.sigma = sigma
        denom = max(float(sigma.max()), float(self.C - sigma.sum()))
        if denom <= 0:
            denom = 1.0
        return (sigma / denom) * self.tau

    def flex_mask(self, weak_probs) -> np.ndarray:
        weak_probs = np.asarray(weak_probs, dtype=np.float64)
        thr = self.update(weak_probs)
        hard = weak_probs.argmax(axis=1)
        conf = weak_probs.max(axis=1)
        return (conf >= thr[hard]).astype(np.float64)

    # TODO: clip(step/length, 0, 1) (handle length<=0 -> 1.0).
    raise NotImplementedError("linear_rampup")

    if length <= 0:
        return 1.0
    return float(np.clip(step / length, 0.0, 1.0))

    # TODO: per the docstring (1.0 once step>=length).
    raise NotImplementedError("cosine_rampup")

    if length <= 0:
        return 1.0
    t = float(np.clip(step / length, 0.0, 1.0))
    return 0.5 * (1.0 - np.cos(np.pi * t))

    # TODO: per the docstring.
    raise NotImplementedError("exp_rampup")

    if length <= 0:
        return 1.0
    t = float(np.clip(step / length, 0.0, 1.0))
    return float(np.exp(-5.0 * (1.0 - t) ** 2))

    # TODO: build the (N,C) complementary one-hot on argmin for confident rows only.
    raise NotImplementedError("complement_distribution")

    weak_probs = np.asarray(weak_probs, dtype=np.float64)
    n, c = weak_probs.shape
    conf = weak_probs.max(axis=1)
    least = weak_probs.argmin(axis=1)
    sel = conf >= tau
    out = np.zeros((n, c), dtype=np.float64)
    rows = np.arange(n)[sel]
    out[rows, least[sel]] = 1.0
    return out

    # TODO: forward value via _stable_log_softmax_np; build out via _mk(val, (logits,)); set out._backward.
    raise NotImplementedError("soft_cross_entropy")

    logp = _stable_log_softmax_np(x)
    row = -(tgt * logp).sum(axis=1)  # (N,)
    val = float((w * row).sum() / n)
    out = _mk(val, (logits,))

    def _bw():
        if not logits.requires_grad:
            return
        sm = np.exp(logp)
        tsum = tgt.sum(axis=1, keepdims=True)  # (N,1)
        grad = (w[:, None] / n) * (sm * tsum - tgt)
        logits._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: forward (sum of squared diffs, weighted, /N); backward through the softmax (see docstring).
    raise NotImplementedError("mse_consistency")

    s = _stable_softmax_np(x)
    diff = s - tgt
    val = float((w * (diff * diff).sum(axis=1)).sum() / n)
    out = _mk(val, (student_logits,))

    def _bw():
        if not student_logits.requires_grad:
            return
        v = (2.0 * w[:, None] / n) * diff
        dot = (v * s).sum(axis=1, keepdims=True)
        grad = s * (v - dot)
        student_logits._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: forward (include the teacher-entropy term so KL is reported); backward as above.
    raise NotImplementedError("kl_divergence")

    logp = _stable_log_softmax_np(x)
    logt = np.log(np.clip(tgt, 1e-12, None))
    row = (tgt * (logt - logp)).sum(axis=1)
    val = float((w * row).sum() / n)
    out = _mk(val, (student_logits,))

    def _bw():
        if not student_logits.requires_grad:
            return
        sm = np.exp(logp)
        tsum = tgt.sum(axis=1, keepdims=True)
        grad = (w[:, None] / n) * (sm * tsum - tgt)
        student_logits._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: forward mean entropy; backward per the docstring.
    raise NotImplementedError("entropy_loss")

    logp = _stable_log_softmax_np(x)
    p = np.exp(logp)
    H = -(p * logp).sum(axis=1)  # (N,)
    val = float(H.sum() / n)
    out = _mk(val, (logits,))

    def _bw():
        if not logits.requires_grad:
            return
        grad = (1.0 / n) * p * (-logp - H[:, None])
        logits._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: forward (-H of the batch marginal); backward per the docstring.
    raise NotImplementedError("class_balance_loss")

    p = _stable_softmax_np(x)
    pbar = p.mean(axis=0)  # (C,)
    val = float((pbar * np.log(np.clip(pbar, 1e-12, None))).sum())
    out = _mk(val, (logits,))

    def _bw():
        if not logits.requires_grad:
            return
        a = 1.0 + np.log(np.clip(pbar, 1e-12, None))  # (C,)
        dot = (a[None, :] * p).sum(axis=1, keepdims=True)  # (N,1)
        grad = (1.0 / n) * p * (a[None, :] - dot)
        logits._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: stable sigmoid-BCE forward via _mk(val, (z,)); backward (w/N)*(sigmoid(z)-t).
    raise NotImplementedError("bce_consistency")

    def _softplus(u):
        return np.logaddexp(0.0, u)
    # -t*log(sigmoid(z)) - (1-t)*log(1-sigmoid(z)) = t*softplus(-z) + (1-t)*softplus(z)
    elem = t * _softplus(-x) + (1.0 - t) * _softplus(x)
    val = float((w * elem.sum(axis=1)).sum() / n)
    out = _mk(val, (z,))

    def _bw():
        if not z.requires_grad:
            return
        s = 1.0 / (1.0 + np.exp(-x))
        grad = (w[:, None] / n) * (s - t)
        z._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: forward sum-squared-diff averaged over N; backward (2w/N)*(x-t).
    raise NotImplementedError("pi_consistency")

    diff = x - t
    val = float((w * (diff * diff).sum(axis=1)).sum() / n)
    out = _mk(val, (z,))

    def _bw():
        if not z.requires_grad:
            return
        grad = (2.0 * w[:, None] / n) * diff
        z._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: forward via _stable_log_softmax_np; backward through softmax per the docstring.
    raise NotImplementedError("focal_consistency")

    logp = _stable_log_softmax_np(x)
    p = np.exp(logp)
    nlog = -logp  # -log p_c
    omp = 1.0 - p
    phi = t * (omp ** gm) * nlog
    val = float((w * phi.sum(axis=1)).sum() / n)
    out = _mk(val, (z,))

    def _bw():
        if not z.requires_grad:
            return
        # u_c = d phi_c / d p_c
        term1 = -gm * (omp ** (gm - 1.0)) * nlog
        term2 = -(omp ** gm) / p
        u = t * (term1 + term2)
        dot = (u * p).sum(axis=1, keepdims=True)
        grad = (w[:, None] / n) * p * (u - dot)
        z._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: forward JS (clip probs to 1e-12 before log); backward through softmax per the docstring.
    raise NotImplementedError("js_divergence")

    p = _stable_softmax_np(x)
    m = 0.5 * (p + t)
    pc = np.clip(p, 1e-12, None)
    tc = np.clip(t, 1e-12, None)
    mc = np.clip(m, 1e-12, None)
    js = 0.5 * (p * (np.log(pc) - np.log(mc))).sum(axis=1) \
        + 0.5 * (t * (np.log(tc) - np.log(mc))).sum(axis=1)
    val = float((w * js).sum() / n)
    out = _mk(val, (z,))

    def _bw():
        if not z.requires_grad:
            return
        u = (w[:, None] / n) * 0.5 * (np.log(pc) - np.log(mc))
        dot = (u * p).sum(axis=1, keepdims=True)
        grad = p * (u - dot)
        z._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: forward -(comp*log(1-p)) (clip 1-p to 1e-12); backward through softmax per the docstring.
    raise NotImplementedError("negative_consistency")

    p = _stable_softmax_np(x)
    omp = np.clip(1.0 - p, 1e-12, None)
    row = -(comp * np.log(omp)).sum(axis=1)
    val = float((w * row).sum() / n)
    out = _mk(val, (z,))

    def _bw():
        if not z.requires_grad:
            return
        u = comp / omp
        dot = (u * p).sum(axis=1, keepdims=True)
        grad = (w[:, None] / n) * p * (u - dot)
        z._accum(out.grad * grad)

    out._backward = _bw
    return out

    # TODO: onehot, mask = pseudo_label(...); loss = soft_cross_entropy(...); scale by lambda_u.
    raise NotImplementedError("consistency_loss")

    weak_probs = np.asarray(weak_probs, dtype=np.float64)
    onehot, mask, hard = pseudo_label(weak_probs, tau)
    loss = soft_cross_entropy(strong_logits, onehot, weight=mask)
    if lambda_u != 1.0:
        loss = ag.mul(loss, float(lambda_u))
    return loss, float(mask.mean())

    # TODO: sup = ag.cross_entropy(...); cons, mr = consistency_loss(...); return (ag.add(sup, cons), mr).
    raise NotImplementedError("fixmatch_loss")

    sup = ag.cross_entropy(sup_logits, sup_y)
    cons, mr = consistency_loss(strong_logits, weak_probs, tau, lambda_u=lambda_u)
    return ag.add(sup, cons), mr

    def update(self, model):
        # TODO: for each named param, shadow <- decay*shadow + (1-decay)*param.data (in place).
        raise NotImplementedError("EMATeacher.update")

    def copy_to(self, model):
        # TODO: write self.shadow[k] into each model param's data.
        raise NotImplementedError("EMATeacher.copy_to")

    def update(self, model):
        d = self.decay
        for k, v in model.named_params().items():
            self.shadow[k] = d * self.shadow[k] + (1.0 - d) * v.data

    def copy_to(self, model):
        for k, v in model.named_params().items():
            v.data[...] = self.shadow[k]

    in_dim = Xl.shape[1]
    C = int(yl.max()) + 1
    # TODO: build the model + optimizer + EMA teacher (+ distribution aligner), then run the
    # FixMatch training loop (supervised CE + masked consistency with confidence-gated, optionally
    # distribution-aligned pseudo-labels), and save_checkpoint(out_path, ema.state_dict(), cfg).
    raise NotImplementedError("train")

    in_dim = Xl.shape[1]
    C = int(yl.max()) + 1
    rng = np.random.default_rng(seed)
    aug = np.random.default_rng(seed + 1234)

    model = MLPClassifier(in_dim, C, d_hidden=D_HIDDEN, n_layer=N_LAYER, seed=seed)
    opt = O.AdamW(model.params(), lr=3e-3, weight_decay=5e-4)
    ema = S.EMATeacher(model, decay=0.999)
    aligner = S.DistributionAligner(C, momentum=0.999)

    Nl = Xl.shape[0]
    Nu = Xu.shape[0]
    lab_bs = min(64, Nl)
    unlab_bs = 128
    base_lr = 3e-3
    warmup = 200
    tau = 0.95
    lambda_u = 1.0
    rampup_len = max(1, steps // 2)

    for step in range(steps):
        lr = lr_at(step, steps, base_lr, warmup)
        opt.lr = lr

        # supervised batch (weak aug)
        li = rng.integers(0, Nl, size=lab_bs)
        Xb = S.weak_augment(Xl[li], aug, sigma=0.1)
        sup = ag.cross_entropy(model(Xb), yl[li])

        # unlabeled batch
        ui = rng.integers(0, Nu, size=unlab_bs)
        Xu_b = Xu[ui]
        wlogits = model(S.weak_augment(Xu_b, aug, sigma=0.1))
        wp = S._stable_softmax_np(np.asarray(wlogits.data, dtype=np.float64))
        wp = aligner.align(wp)
        slogits = model(S.strong_augment(Xu_b, aug, sigma=0.5, mask_frac=0.5))
        w_unl = lambda_u * S.linear_rampup(step, rampup_len)
        cons, mr = S.consistency_loss(slogits, wp, tau=tau, lambda_u=w_unl)

        loss = ag.add(sup, cons)
        opt.zero_grad()
        loss.backward()
        O.clip_grad_norm(model.params(), 5.0)
        opt.step()
        em
…[truncated 417 chars]