Catch bugs locally...

def mtf_transformer_tiny():
  """Catch bugs locally..."""
  hparams = mtf_transformer_base()
  hparams.d_model = 128
  hparams.d_ff = 512
  hparams.batch_size = 8
  hparams.encoder_layers = ["att", "drd"] * 2
  hparams.decoder_layers = ["att", "enc_att", "drd"] * 2
  hparams.num_heads = 8
  # data parallelism and model-parallelism
  hparams.mesh_shape = "batch:2;model:4"
  hparams.activation_dtype = "float32"
  return hparams