Sync model.py for 263702

model.py

@@ -173,6 +173,20 @@ def soft_targets(minutes: torch.Tensor, bins: int, kappa: float = 40.0):
     return out
 
 
+def ring_target(minutes: torch.Tensor, bins: int, kappa: float = 40.0):
+    """A von Mises bump over the whole 12-hour ring, for the time-itself head.
+
+    The angular heads are trained on hand directions; this one is trained on the
+    time, so bin b covers minutes b/bins of the way round 720 rather than round
+    a single revolution of a hand.
+    """
+    dev = minutes.device
+    centres = (torch.arange(bins, device=dev, dtype=torch.float32) + 0.5) / bins * 2 * math.pi
+    d = centres.unsqueeze(0) - (minutes / 720.0 * 2 * math.pi).unsqueeze(1)
+    t = torch.exp(kappa * (torch.cos(d) - 1.0))
+    return t / t.sum(dim=1, keepdim=True)
+
+
 def soft_argmax_angle(logits: torch.Tensor):
     """Circular expectation of a distribution over angle bins -> radians."""
     bins = logits.shape[1]