GPU and batching

class ResidualBlock(nn.Module):
    def __init__(self, ch: int, norm: str = "group") -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
        self.bn1 = make_norm(norm, ch)
        self.conv2 = nn.Conv2d(ch, ch, 3, padding=1, bias=False)
        self.bn2 = make_norm(norm, ch)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = F.relu(self.bn1(self.conv1(x)), inplace=True)
        y = self.bn2(self.conv2(y))
        return F.relu(x + y, inplace=True)


class Drop7Net(nn.Module):
    """Small residual CNN with a 7-way column policy head and a scalar head.

    Input:  (N, 12, 7, 7)
    Output: (N, 7) column logits, (N,) scalar regression (e.g. lifetime score).
    """

    def __init__(self, channels: int = 128, blocks: int = 6,
                 in_channels: int = IN_CHANNELS, norm: str = "group") -> None:
        super().__init__()
        self.norm_kind = norm
        self.stem = nn.Sequential(
            nn.Conv2d(in_channels, channels, 3, padding=1, bias=False),
            make_norm(norm, channels),
            nn.ReLU(inplace=True),
        )
        self.blocks = nn.Sequential(
            *[ResidualBlock(channels, norm) for _ in range(blocks)])
        # Policy head: 7 column logits.
        self.policy = nn.Sequential(
            nn.Conv2d(channels, 32, 1, bias=False),
            make_norm(norm, 32), nn.ReLU(inplace=True),
            nn.Flatten(), nn.Linear(32 * BOARD * BOARD, N_COLUMNS),
        )
        # Scalar head: one regression output.
        self.value = nn.Sequential(
            nn.Conv2d(channels, 32, 1, bias=False),
            make_norm(norm, 32), nn.ReLU(inplace=True),
            nn.Flatten(), nn.Linear(32 * BOARD * BOARD, 128),
            nn.ReLU(inplace=True), nn.Linear(128, 1),
        )

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        h = self.blocks(self.stem(x))
        return self.policy(h), self.value(h).squeeze(-1)

def make_norm(kind: str, ch: int) -> nn.Module:
    """Norm layer factory.

    'batch' is the textbook choice but is BROKEN FOR TRAINING on gfx1151 with
    the MIOpen shipped inside the torch ROCm wheel: MIOpen selects a GFX9-only
    solver for the batch-norm training kernel and fails to assemble it
    ("v_add_f32 ... row_bcast:31" is not a valid gfx1151 operand), so every
    training step raises miopenStatusUnknownError. Eval mode is unaffected.

    'group' avoids MIOpen entirely -- GroupNorm is a native PyTorch kernel --
    and is the default here. It is also the better choice on the merits for a
    self-play/RL setting, where BatchNorm's running statistics drift between
    the acting and learning distributions and couple samples within a batch.
    """
    if kind == "batch":
        return nn.BatchNorm2d(ch)
    if kind == "group":
        return nn.GroupNorm(min(8, ch), ch)
    if kind == "none":
        return nn.Identity()
    raise ValueError(f"unknown norm kind {kind!r}")

# --------------------------------------------------------------------------
# 2. Correctness
# --------------------------------------------------------------------------

def correctness(dev: torch.device, tol_matmul: float = 5e-2) -> dict[str, Any]:
    """A GPU that imports but computes garbage is a real failure mode on
    partially supported targets. These checks compare against CPU."""
    section("2. NUMERICAL CORRECTNESS (GPU vs CPU)")
    res: dict[str, Any] = {"passed": True, "checks": []}

    def record(name: str, ok: bool, detail: str, gate: bool = True) -> None:
        """gate=False: a known, documented limitation. Still reported loudly,
        but it does not fail the run, because the supported configuration
        works around it (see make_norm)."""
        res["checks"].append({"name": name, "ok": bool(ok),
                              "detail": detail, "gate": gate})
        if gate:
            res["passed"] = res["passed"] and bool(ok)
        else: