code: def batch_norm_relu(inputs, is_training, relu=True):
    
    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
docs: Block of batch norm and relu.