code: def __init__(
      self,
      dim,
      num_experts=(4, 4),
      hidden_dim=None,
      activation=nn.ReLU,
      second_policy_train="random",
      second_policy_eval="random",
      second_threshold_train=0.2,
      second_threshold_eval=0.2,
      capacity_factor_train=1.25,
      capacity_factor_eval=2.0,
      loss_coef=1e-2,
      experts=None,
  ):
      super().__init__()
  
      assert (
          len(num_experts) == 2
      ), "only 2 levels of heirarchy for experts allowed for now"
      num_experts_outer, num_experts_inner = num_experts
      self.num_experts_outer = num_experts_outer
      self.num_experts_inner = num_experts_inner
  
      gating_kwargs = {
          "second_policy_train": second_policy_train,
          "second_policy_eval": second_policy_eval,
          "second_threshold_train": second_threshold_train,
          "second_threshold_eval": second_threshold_eval,
          "capacity_factor_train": capacity_factor_train,
          "capacity_factor_eval": capacity_factor_eval,
      }
  
      self.gate_outer = Top2Gating(
          dim, num_gates=num_experts_outer, **gating_kwargs
      )
      self.gate_inner = Top2Gating(
          dim,
          num_gates=num_experts_inner,
          outer_expert_dims=(num_experts_outer,),
          **gating_kwargs,
      )
  
      self.experts = default(
          experts,
          lambda: Experts(
              dim,
              num_experts=num_experts,
              hidden_dim=hidden_dim,
              activation=activation,
          ),
      )
      self.loss_coef = loss_coef
  
imports: import torch
  from torch import nn
  import torch.nn.functional
  import math
  from inspect import isfunction
label_desc: Use image node_id 1 to create a new HeirarchicalMoE object from inherited base classes: nn with example: obj = HeirarchicalMoE(dim, num_experts, hidden_dim, activation, second_policy_train, second_policy_eval, second_threshold_train, second_threshold_eval, capacity_factor_train, capacity_factor_eval, loss_coef, experts)