code: def auto_batch_size(sequence_length,
                      mesh_shape,
                      layout_rules,
                      tokens_per_split=2048):
    
    num_splits = mtf.tensor_dim_to_mesh_dim_size(
        layout_rules, mesh_shape, mtf.Dimension("batch", 0))
    ret = max(1, tokens_per_split // sequence_length) * num_splits
    tf.logging.info(
        "AUTO_BATCH_SIZE tokens_per_split=%s num_splits=%s"
        " sequence_length=%s batch_size=%s"
        % (tokens_per_split, num_splits, sequence_length, ret))
    return ret
docs: Automatically compute batch size.
  
    Args:
      sequence_length: an integer
      mesh_shape: an input to mtf.convert_to_shape()
      layout_rules: an input to mtf.convert_to_layout_rules()
      tokens_per_split: an integer
    Returns:
      an integer