"""Shared helpers for backbones."""

from __future__ import annotations

import torch


def _pack_scalar(s: torch.Tensor, B: int) -> torch.Tensor:
    """Normalize a scalar / per-batch scalar to shape (B, 1)."""
    if s.ndim == 0 or (s.ndim == 1 and s.numel() == 1):
        return s.view(1, 1).expand(B, 1)
    return s.reshape(B, 1)


def count_parameters(model: torch.nn.Module) -> int:
    return sum(p.numel() for p in model.parameters() if p.requires_grad)