def xmoe_2d_x64():
  """Two-dimensional hierarchical mixture of 64 experts."""
  hparams = xmoe_2d()
  # hparams.mesh_shape = "b0:4;b1:8"
  hparams.outer_batch_size = 4
  hparams.moe_num_experts = [8, 8]
  return hparams
Two-dimensional hierarchical mixture of 64 experts.
tensor2tensor/models/research/moe_experiments.py