Project: 073_ChartMoE_Mixture_of_Diversely_Aligned_Expert_Conne | Function: forward (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Enhanced MoE routing with load balancing losses including router z-loss and balance loss computation for training stability

**KEY TERMS:** {
  "gate_logits": "gj(vi)",
  "gate_softmax": "σ(MG(vi))",
  "selected_experts": "Top(·; K)",
  "weights": "gj(vi)",
  "x_img": "vi"
}

**CONTEXT CODE:** """
    FEATURE: Implementation of MLP-MoE
    AUTHOR: Brian Qu
    URL: https://arxiv.org/abs/2409.03277
    REFERENCE: https://github.com/SHI-Labs/CuMo
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat, reduce, pack, unpack

class MLPMoE(nn.Module):
    def __init__(self, num_experts, num_selected, mm_channels, channels):
        super().__init__()
        self.num_experts = num_experts
        self.num_selected = num_selected
        self.mm_channels = mm_channels
        self.channels = channels

        self.gate = nn.Linear(mm_channels, num_experts, bias=False)
        # nn.init.zeros_(self.gate.weight)

        self.num_selected = num_selected
        self.num_experts = num_experts
        
        self.experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(mm_channels, channels, bias=True), 
                    nn.GELU(), 
                    nn.Linear(channels, channels, bias=True)
                )
                for _ in range(num_experts)
            ]
        )
        
    def forward(self, x_img):
        raise NotImplementedError("Function implementation has been masked for testing")
    
class MLPMoE_bzloss(nn.Module):
    def __init__(self, num_experts, num_selected, mm_channels, channels):
        super().__init__()
        self.num_experts = num_experts
        self.num_selected = num_selected
        self.mm_channels = mm_channels
        self.channels = channels

        self.gate = nn.Linear(mm_channels, num_experts, bias=False)
        # nn.init.zeros_(self.gate.weight)

        self.num_selected = num_selected
        self.num_experts = num_experts
        self.experts = nn.ModuleList([nn.Sequential(nn.Linear(mm_channels, channels, bias=True), nn.GELU(), nn.Linear(channels, channels, bias=True)) for _ in range(num_experts)])

    def forward(self, x_img):
        raise NotImplementedError("Function implementation has been masked for testing")

Complete the function implementation replacing "raise NotImplementedError".