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