docstr: x: (sequence_length, model_dim)
  gate_logits: (sequence_length, n_experts)
docstr_tok: ["x",":","(","sequence_length",",","model_dim",")","gate_logits",":","(","sequence_length",",","n_experts",")"]
code: def forward(self, x: torch.Tensor) -> torch.Tensor:
      """
      x: (sequence_length, model_dim)
      gate_logits: (sequence_length, n_experts)
      """
      # optional reshape
      input_shape = x.shape
      x = x.view(-1, input_shape[-1])
  
      # gate_logits: (sequence_length, n_experts)
      gate_logits = self.gate(x)
      # all_probs: (sequence_length, n_experts) and upcast for softmax
      all_probs = torch.nn.functional.softmax(
          gate_logits, dim=1, dtype=torch.float
      )
  
      if self.top_k < self.num_experts:
          _, not_selected_experts = torch.topk(
              all_probs,
              self.num_experts - self.top_k,
              largest=False,
              sorted=False,
              dim=1,
          )
          # Mask not selected experts
          all_probs.scatter_(1, not_selected_experts, 0)
  
      # Re-normalize
      weights = all_probs / all_probs.sum(dim=1, keepdim=True)
  
      # Final output tensor
      out = x.new_zeros(x.shape[0], self.hidden_dim)
      for i in range(self.num_experts):
          h = self.act(self.w1[i](x)) * self.w3[i](x)
          h = self.w2[i](h, reduce=False)
          # Add expert output to out with masking
          out += h * weights[:, i].view(-1, 1)
  
      # Reduce sum
      if self.process_group.size() > 1:
          torch.distributed.all_reduce(out, group=self.process_group)
  
      return out
  
imports: import torch
  import torch.distributed
  import numpy
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.configuration_utils import PretrainedConfig
  from typing import Optional, List, Tuple
  from loguru import logger
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.layers import FastLinear, FastRMSNorm, TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, PositionRotaryEmbedding, TensorParallelHead, get_linear
label_desc: Use image node_id 2 for calling the DenseMoE obj's underlying member method code with example usage: obj.forward(x) and returns: out