Multi-Faceted Modeling Approach Questions

import torch
import torch.nn as nn
import torch.nn.functional as F

class GatingNetwork(nn.Module):
    def __init__(self, in_dim, hidden=32, n_experts=3):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden)
        self.fc2 = nn.Linear(hidden, n_experts)  # outputs logits for experts

    def forward(self, x):
        x = F.relu(self.fc1(x))
        logits = self.fc2(x)            # shape: (batch, n_experts)
        probs = F.softmax(logits, dim=-1)
        return logits, probs

# Example expert stubs (could be full models)
class Expert(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.net = nn.Linear(in_dim, out_dim)
    def forward(self, x):
        return self.net(x)

# Instantiate
in_dim = 16
out_dim = 4
gate = GatingNetwork(in_dim)
experts = [Expert(in_dim, out_dim) for _ in range(3)]
fallback_expert = Expert(in_dim, out_dim)  # shared fallback

# Inference-time routing with fallback
def route_inference(x, gate, experts, fallback_expert, available_mask=None):
    # x: (batch, in_dim); available_mask: list/tensor of booleans length n_experts
    _, probs = gate(x)
    top1_idx = torch.argmax(probs, dim=-1)  # (batch,)
    results = []
    n = len(experts)
    if available_mask is None:
        available_mask = torch.ones(n, dtype=torch.bool)

    for i, xi in enumerate(x):
        chosen = int(top1_idx[i].item())
        # if chosen unavailable, pick next-best available
        if not available_mask[chosen]:
            # sort indices by descending prob, find first available
            sorted_idxs = torch.argsort(probs[i], descending=True)
            found = False
            for idx in sorted_idxs:
                if available_mask[int(idx.item())]:
                    chosen = int(idx.item())
                    found = True
                    break
            if not found:
                # all experts unavailable -> fallback_expert
                out = fallback_expert(xi.unsqueeze(0))
                results.append(out)
                continue
        out = experts[chosen](xi.unsqueeze(0))
        results.append(out)
    return torch.cat(results, dim=0)  # (batch, out_dim)