code: def generate_heterogeneous_expert_masks(
      mask_info, num_experts, experts_dim, mesh, expert_width
  ):
      """Generates the heterogeous expert masks.
  
      Example mask_info format:
        mask_info = [{'percent_number': .5, 'layers': 1, 'width':1},
                     {'percent_number': .5, 'layers': 2, 'width':2}]
  
      Args:
        mask_info: list of dicts.
        num_experts: number of experts in the model
        experts_dim: mtf dimension for experts (partitioned)
        mesh: mesh object
        expert_width: int, default expert width which will be modified by the mask
  
      Returns:
        mask of shape [moe_num_layers, num_experts, hidden_size].
      """
      # Get max num layers
      max_layers = max([m["layers"] for m in mask_info])
      # Get max width
      max_width = max([m["width"] for m in mask_info])
      # Will be shape [max_width, max_layers, num_experts]
      expert_mask = np.zeros([max_width, max_layers, 0])
      for idx, mask_i in enumerate(mask_info):
          if mask_i["percent_number"] < 1.0:
              num_experts_in_mask = int(
                  num_experts * mask_i["percent_number"]
              )
          else:
              num_experts_in_mask = int(mask_i["percent_number"])
          # if percent_number=1 either because homogeneous experts or just 1 expert
          # in which case num_experts_in_mask will be reset to num_experts
          # creating one homogeneous group
          if idx == (len(mask_info) - 1):  # Last position
              num_experts_in_mask_tmp = (
                  num_experts - expert_mask.shape[2]
              )
              if num_experts_in_mask_tmp != num_experts_in_mask:
                  tf.logging.info(
                      "Expert layer probabilities do not evenly divide "
                      "the number of experts: {} {}".format(
                          num_experts_in_mask, num_experts_in_mask_tmp
                      )
                  )
                  num_experts_in_mask = num_experts_in_mask_tmp
          mask = np.zeros(
              [int(max_width), int(max_layers), num_experts_in_mask]
          )
          # Zero out the last layers of the experts.
          mask[
              : (mask_i["width"] * expert_width), : mask_i["layers"], :
          ] = 1
          expert_mask = np.concatenate(
              [expert_mask, mask], axis=2
          )  # expert dim
      assert expert_mask.shape[2] == num_experts
      tf.logging.info("heterogeneous mask: {}".format(expert_mask))
  
      # Now import the numpy mask into Mesh TF.
      layers_dim = mtf.Dimension("num_expert_layers", max_layers)
      width_dim = mtf.Dimension("expert_hidden", max_width)
      expert_mask_tf = tf.convert_to_tensor(expert_mask)
      expert_mask_mtf = mtf.import_tf_tensor(
          mesh,
          tf_tensor=expert_mask_tf,
          shape=[width_dim, layers_dim, experts_dim],
      )
      return expert_mask_mtf
  
file_path: mesh/mesh_tensorflow/transformer/heterogeneous_moe.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import moe
  from mesh_tensorflow.transformer import transformer
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: generate_heterogeneous_expert_masks(mask_info, num_experts, experts_dim, mesh, expert_width) and returns: expert_mask_mtf