>>109706688
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class UltrametricLoadoutRouter(nn.Module):
def __init__(self, feature_dim, depth=4, arity=4):
super().__init__()
self.L, self.p = depth, arity
# Maps raw piece stats / team embeddings to tree logits
self.proj = nn.Linear(feature_dim, depth * arity)
def forward(self, loadout_features, tau=1.0, hard=False):
# Shape: (Batch, L, p)
logits = self.proj(loadout_features).view(-1, self.L, self.p)
# Soft / Gumbel categorical routing distribution at each depth
probs = F.gumbel_softmax(logits, tau=tau, hard=hard)
# Branch overlap at each depth level l between loadout i and loadout j
# M[i, j, l] = sum_c (prob_i[l, c] * prob_j[l, c])
M = torch.einsum('ilc, jlc -> ijl', probs, probs)
# Continuous lowest common ancestor depth via cumulative minimum
# (Once branches diverge at level l, all deeper levels are cut)
branch_continuity = torch.cummin(M, dim=-1).values
d_p = self.L - branch_continuity.sum(dim=-1)
return probs, d_p
