Rank-2 tensors for all muon-wise variables

[GilesStrong]

Currently code uses a mixture of rank-1 and rank-2 tensors for muon-wise variables, e.g. the momentum of a batch of muons has shape [300], but the theta_msc of their trajectories has shape [300,1]. Devs must then be careful to squeeze or unsqueeze tensors in calculations, since the rank-1 tensors can be accidentally broadcast across the rank-2 tensors, e.g. AbsX0Inferer._x0_from_dtheta can return a [300,300] tensor instead of a [300] or [300,1] tensor if the momentum is not unsqueezed.
I propose that all muon-wise variables use at least rank-2 tensors, even if the length of the last dimension is just one. This should help prevent accidental broadcasting.