code: def tensor_dim_to_mesh_dim_size(layout, mesh_shape, tensor_dim):
    
    layout_rules = convert_to_layout_rules(layout)
    mesh_shape = convert_to_shape(mesh_shape)
    mesh_axis = layout_rules.tensor_dimension_to_mesh_axis(tensor_dim, mesh_shape)
    if mesh_axis is None:
      return 1
    else:
      return mesh_shape.dims[mesh_axis].size
docs: How many ways does a tensor dimension get split.
  
    This is used to "cheat" when building the mtf graph and peek at how a
    tensor dimension will be split.  Returns 1 if the tensor dimension is not
    split.
  
    Args:
      layout: an input to convert_to_layout_rules
      mesh_shape: an input to convert_to_shape
      tensor_dim: a Dimension
  
    Returns:
      an integer