Random parameters

class Choice:
    """A discrete set of values to sample from.

    Args:
        values: Sequence of numeric values to choose from.
        probabilities: Optional per-value probabilities.
            If `None`, all values are equally likely.

    Examples:
        >>> Choice([-10, 0, 10])
        Choice([-10.0, 0.0, 10.0])
        >>> Choice([0.5, 1.0, 2.0], probabilities=[0.2, 0.6, 0.2])
        Choice([0.5, 1.0, 2.0], p=[0.20, 0.60, 0.20])
    """

    def __init__(
        self,
        values: Sequence[float | int],
        probabilities: Sequence[float] | None = None,
    ) -> None:
        if len(values) < 1:
            msg = "Choice requires at least one value"
            raise ValueError(msg)
        self._values = torch.tensor([float(v) for v in values])
        if probabilities is not None:
            if len(probabilities) != len(values):
                msg = f"Expected {len(values)} probabilities, got {len(probabilities)}"
                raise ValueError(msg)
            self._probs = torch.tensor([float(p) for p in probabilities])
        else:
            self._probs = torch.ones(len(values)) / len(values)

    def sample(self) -> float:
        """Pick one value at random."""
        idx = torch.multinomial(self._probs, 1).item()
        return float(self._values[int(idx)])

    def sample_batched(self, n: int) -> torch.Tensor:
        """Pick ``n`` values at random, with replacement.

        Args:
            n: Number of independent draws.

        Returns:
            A 1D tensor of shape ``(n,)``.
        """
        idx = torch.multinomial(self._probs, n, replacement=True)
        return self._values[idx]

    def __repr__(self) -> str:
        vals = [f"{v:.1f}" if v == int(v) else f"{v}" for v in self._values.tolist()]
        parts = f"[{', '.join(vals)}]"
        if torch.allclose(self._probs, self._probs[0].expand_as(self._probs)):
            return f"Choice({parts})"
        probs = ", ".join(f"{p:.2f}" for p in self._probs.tolist())
        return f"Choice({parts}, p=[{probs}])"