def mtf_image_transformer_tiny():
  """Catch bugs locally..."""
  hparams = mtf_image_transformer_base()
  hparams.hidden_size = 128
  hparams.d_ff = 256
  hparams.batch_size = 4
  hparams.num_encoder_layers = 1
  hparams.num_decoder_layers = 4
  hparams.num_heads = 4
  hparams.attention_key_size = 128
  hparams.attention_value_size = 128
  hparams.block_length = 32
  # data parallelism and model-parallelism
  hparams.mesh_shape = "batch:2"
  hparams.layout = "batch:batch"
  return hparams
Catch bugs locally...