Block of batch norm and relu.

def batch_norm_relu(inputs, is_training, relu=True):
  """Block of batch norm and relu."""
  inputs = mtf.layers.batch_norm(
      inputs,
      is_training,
      BATCH_NORM_DECAY,
      epsilon=BATCH_NORM_EPSILON,
      init_zero=(not relu))
  if relu:
    inputs = mtf.relu(inputs)
  return inputs