Transformer with mixture of experts.  890M Params.

def lmx_moe_h1k_f4k_x32():
  """Transformer with mixture of experts.  890M Params."""
  hparams = lmx_h1k_f4k()
  hparams.ffn_layer = "local_moe_tpu"
  hparams.moe_num_experts = 32
  hparams.weight_dtype = "bfloat16"
  hparams.batch_size = 8192
  return hparams