Spaces:
Sleeping
Sleeping
| """Alternative point estimators for axial categorical distributions.""" | |
| import torch | |
| def estimates(logits, names): | |
| probability = logits.float().softmax(-2) | |
| height = logits.shape[-2] | |
| grid = torch.arange(height, device=logits.device, dtype=torch.float32)[ | |
| None, None, :, None | |
| ] | |
| result = {} | |
| if "mean" in names: | |
| result["mean"] = (probability * grid).sum(-2) | |
| if "median" in names: | |
| cumulative = probability.cumsum(-2) | |
| index = (cumulative < 0.5).sum(-2).clamp_max(height - 1) | |
| mass = probability.gather(-2, index.unsqueeze(-2)).squeeze(-2) | |
| before = cumulative.gather(-2, (index - 1).clamp_min(0).unsqueeze(-2)).squeeze( | |
| -2 | |
| ) | |
| before = torch.where(index == 0, torch.zeros_like(before), before) | |
| result["median"] = ( | |
| index.float() - 0.5 + (0.5 - before) / mass.clamp_min(1e-12) | |
| ).clamp(0, height - 1) | |
| if "quadratic_mode" in names: | |
| index = logits.argmax(-2) | |
| score = logits.float() | |
| center = score.gather(-2, index.unsqueeze(-2)).squeeze(-2) | |
| left = score.gather(-2, (index - 1).clamp_min(0).unsqueeze(-2)).squeeze(-2) | |
| right = score.gather( | |
| -2, (index + 1).clamp_max(height - 1).unsqueeze(-2) | |
| ).squeeze(-2) | |
| curvature = left - 2 * center + right | |
| offset = torch.where( | |
| (curvature < -1e-8) & (index > 0) & (index < height - 1), | |
| 0.5 * (left - right) / curvature.clamp_max(-1e-8), | |
| torch.zeros_like(center), | |
| ).clamp(-0.5, 0.5) | |
| result["quadratic_mode"] = (index.float() + offset).clamp(0, height - 1) | |
| if set(result) != set(names): | |
| raise ValueError("Unknown posterior estimator") | |
| return result | |
| def positions_from_logits(logits, config, device=None): | |
| by_device = config.get("inference", {}).get("posterior_by_device") | |
| if by_device is not None: | |
| if device is None or device not in by_device: | |
| raise ValueError("Configured head estimators require the device") | |
| keys = config["device_keys"][device] | |
| mapping = by_device[device] | |
| if set(mapping) != set(keys) or logits.shape[1] != len(keys): | |
| raise ValueError("Configured head estimators differ from output identities") | |
| names = [mapping[key] for key in keys] | |
| values = estimates(logits, sorted(set(names))) | |
| return torch.stack( | |
| [values[name][:, index] for index, name in enumerate(names)], dim=1 | |
| ) | |
| name = config.get("inference", {}).get("posterior", "mean") | |
| if name not in ["mean", "median", "quadratic_mode"]: | |
| raise ValueError("Unsupported configured posterior") | |
| return estimates(logits, [name])[name] | |