def xmoe2_tiny():
  """Test on local cpu."""
  hparams = xmoe2_v1()
  hparams.decoder_layers = [
      "local_att", "att", "compressed_att", "drd", "hmoe"]
  hparams.d_model = 128
  hparams.moe_hidden_size = 512
  hparams.outer_batch_size = 0
  hparams.batch_size = 2
  hparams.mesh_shape = ""
  hparams.activation_dtype = "float32"
  return hparams
Test on local cpu.