"""The packed-batch record the model reads (a copy of model/loader.py's Batch, without the training loader)."""

from __future__ import annotations

from dataclasses import dataclass

from torch import Tensor


@dataclass
class Batch:
    """One packed microbatch on the GPU: every sequence is its own segment, positions restart at 0."""

    x: Tensor  # (N,) int64 inputs
    y: Tensor  # (N,) int64 targets (ignored by inference)
    pos: Tensor  # (N,) int32 position inside its segment (RoPE)
    cu: Tensor  # (S+1,) int32 segment offsets
    cu64: Tensor  # same offsets, int64 (fla)
    cu_host: Tensor  # same offsets, int64 on the CPU
    row: int  # longest segment
    segs: int  # segment count
    weight: Tensor | None = None  # (N,) loss weight per target; None = plain mean

    @property
    def tokens(self) -> int:
        return self.x.numel()