I would like to alter the model below to dynamically allocate experts to a designated cluster. I would like there to be a constant number of experts and clusters but the number of experts per cluster can be the changeable variable. The objective is for the model to arrive to the best combination of experts per cluster that achieves the best results. An inspiration for this could be similar to how a softmax function distributes probabilities across various categories with the added probabilities equalling one. Let me know if you need further clarification with some of the ideas! Here is the model:
```
# Define hyperparameters
batch_size = 64  # Number of samples processed before the model is updated
block_size = 256  # Length of sequences processed by the model
max_iters = 5001  # Maximum number of training iterations
eval_interval = 500  # Interval at which the model is evaluated
learning_rate = 3e-4  # Learning rate for the optimizer
device = 'cuda' if torch.cuda.is_available() else 'cpu'  # Use GPU if available
eval_iters = 200  # Number of iterations for evaluation
n_embd = 384  # Dimension of the embeddings
n_head = 6  # Number of attention heads
n_layer = 6  # Number of transformer layers
dropout = 0.2  # Dropout rate to prevent overfitting
n_clusters = 4  # Number of clusters in the MoE layer
experts_per_cluster = 2  # Number of experts per cluster in the MoE layer
weight_decay = 0.1  # Weight decay for regularization
aux_loss_weight = 0.01  # Weight of the auxiliary loss from the MoE gating mechanism

# Set random seed for reproducibility
torch.manual_seed(1337)

# Define the ClusteredSwitchFeedForward class
class ClusteredSwitchFeedForward(nn.Module):
    """Clustered Mixture-of-Experts (MoE) FeedForward layer that activates all experts within a cluster."""

    def __init__(self):
        super().__init__()
        self.n_clusters = n_clusters
        self.experts_per_cluster = experts_per_cluster

        # Initialize experts for each cluster
        self.experts = nn.ModuleList([
            nn.ModuleList([
                nn.Sequential(
                    nn.Linear(n_embd, 4 * n_embd),
                    nn.ReLU(),
                    nn.Linear(4 * n_embd, n_embd),
                    nn.Dropout(dropout)
                ) for _ in range(self.experts_per_cluster)
            ]) for _ in range(self.n_clusters)
        ])

        # Gating networks
        self.cluster_gate = nn.Linear(n_embd, self.n_clusters)
        self.expert_gate = nn.ModuleList([
            nn.Linear(n_embd, self.experts_per_cluster) for _ in range(self.n_clusters)
        ])

    def forward(self, x):
        """
        Args:
            x: Tensor of shape (B, T, D)

        Returns:
            output: Tensor of shape (B, T, D) after processing through all experts
            total_aux_loss: Scalar tensor representing the sum of auxiliary losses
        """
        B, T, D = x.shape

        # Step 1: Cluster Selection
        cluster_logits = self.cluster_gate(x)  # (B, T, n_clusters)
        cluster_probs = F.softmax(cluster_logits, dim=-1)  # (B, T, n_clusters)
        cluster_pred = torch.argmax(cluster_probs, dim=-1)  # (B, T)

        # Auxiliary loss for cluster gating
        cluster_loss_aux = F.cross_entropy(
            cluster_logits.view(-1, cluster_logits.size(-1)),
            cluster_pred.view(-1)
        )

        # Initialize output tensor
        output = torch.zeros_like(x)  # (B, T, D)
        expert_loss_aux = 0.0

        # Step 2: Expert Processing within Clusters
        for cluster_idx in range(self.n_clusters):
            cluster_mask = cluster_pred == cluster_idx  # (B, T)
            if cluster_mask.any():
                # Select tokens assigned to the current cluster
                cluster_tokens = x[cluster_mask]  # (N, D)

                # Expert gating within the cluster
                expert_logits = self.expert_gate[cluster_idx](cluster_tokens)  # (N, experts_per_cluster)
                expert_probs = F.softmax(expert_logits, dim=-1)  # (N, experts_per_cluster)

                # Auxiliary loss for expert gating
                expert_loss_aux += F.cross_entropy(
                    expert_logits.view(-1, expert_logits.size(-1)),
                    torch.argmax(expert_logits, dim=-1).view(-1)
                )

                # Process with all experts and aggregate outputs
                cluster_output = torch.zeros_like(cluster_tokens)  # (N, D)
                for expert_idx, expert in enumerate(self.experts[cluster_idx]):
                    expert_output = expert(cluster_tokens)  # (N, D)
                    # Weight expert outputs by gating probabilities
                    cluster_output += expert_probs[:, expert_idx].unsqueeze(-1) * expert_output  # (N, D)

                # Assign aggregated outputs back to the main output tensor
                output[cluster_mask] = cluster_output  # (B, T, D)

        # Step 3: Combine Auxiliary Losses
        total_aux_loss = cluster_loss_aux + expert_loss_aux

        return output, total_aux_loss

# Define the Head class
class Head(nn.Module):
    """One head of self-attention."""

    def __init__(self, head_size):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)    # (B, T, head_size)
        q = self.query(x)  # (B, T, head_size)
        wei = q @ k.transpose(-2, -1) * C**-0.5  # (B, T, T)
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))  # Apply causal mask
        wei = F.softmax(wei, dim=-1)  # (B, T, T)
        wei = self.dropout(wei)
        v = self.value(x)  # (B, T, head_size)
        out = wei @ v      # (B, T, head_size)
        return out

# Define the MultiHeadAttention class
class MultiHeadAttention(nn.Module):
    """Multiple heads of self-attention in parallel."""

    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(head_size * num_heads, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)  # Concatenate all heads
        out = self.dropout(self.proj(out))
        return out

# Define the Block_ClusteredSwitch class
class Block_ClusteredSwitch(nn.Module):
    """Transformer block with Clustered Mixture-of-Experts FeedForward layer."""

    def __init__(self):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
        self.sa = MultiHeadAttention(n_head, head_size=n_embd // n_head)
        self.ffwd = ClusteredSwitchFeedForward()

    def forward(self, x):
        x = x + self.sa(self.ln1(x))  # Self-Attention with residual connection
        ffwd_out, loss_aux = self.ffwd(self.ln2(x))  # Clustered MoE FeedForward with auxiliary loss
        x = x + ffwd_out  # FeedForward with residual connection
        return x, loss_aux

# Define the Clustered_ST class
class Clustered_ST(nn.Module):
    """Clustered Switch Transformer with Clustered Mixture-of-Experts."""

    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([Block_ClusteredSwitch() for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)  # Final layer norm
        self.lm_head = nn.Linear(n_embd, vocab_size)

        self.apply(self._init_weights)

    def _init_weights(self, module):
        """Initialize weights with normal distribution."""
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding_table(idx)  # (B, T, n_embd)
        pos_emb = self.position_embedding_table(torch.arange(T, device=device))  # (T, n_embd)
        x = tok_emb + pos_emb  # (B, T, n_embd)
        aux_loss = 0.0
        for block in self.blocks:
            x, loss_aux = block(x)
            aux_loss += loss_aux  # Accumulate auxiliary loss
        x = self.ln_f(x)  # (B, T, n_embd)
        logits = self.lm_head(x)  # (B, T, vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets) + aux_loss_weight * aux_loss  # Total loss with auxiliary term

        return logits, loss

    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]  # Focus on the last time step
            probs = F.softmax(logits, dim=-1)  # Convert to probabilities
            idx_next = torch.multinomial(probs, num_samples=1)  # Sample the next token
            idx = torch.cat((idx, idx_next), dim=1)  # Append to the sequence
        return idx
```