The number of entries in a tensor.

    If partial_layout is specified, then mesh_dimension_to_size must also be. In
    this case, the number of entries on a single device is returned.

    Args:
      tensor_name: a string, name of a tensor in the graph.
      partial_layout: an optional {string: string}, from MTF dimension name to
          mesh dimension name.
      mesh_dimension_to_size: an optional {string: int}, from mesh dimension
          name to size.

    Returns:
      an integer

def get_tensor_num_entries(self, tensor_name, partial_layout=None,
                             mesh_dimension_to_size=None):
    """The number of entries in a tensor.

    If partial_layout is specified, then mesh_dimension_to_size must also be. In
    this case, the number of entries on a single device is returned.

    Args:
      tensor_name: a string, name of a tensor in the graph.
      partial_layout: an optional {string: string}, from MTF dimension name to
          mesh dimension name.
      mesh_dimension_to_size: an optional {string: int}, from mesh dimension
          name to size.

    Returns:
      an integer
    """
    shape = self.get_tensor_shape(tensor_name)
    # We don't have to worry about divisiblity issues because Mesh TensorFlow
    # only allows evenly divisible assignments.
    num_entries = 1
    for dim in shape.dims:
      num_entries = num_entries * dim.value

    if not partial_layout:
      return num_entries

    for mtf_dimension_name in self.get_tensor_mtf_dimension_names(tensor_name):
      if mtf_dimension_name not in partial_layout:
        continue
      mesh_dimension_name = partial_layout[mtf_dimension_name]
      mesh_dimension_size = mesh_dimension_to_size[mesh_dimension_name]
      num_entries = int(math.ceil(num_entries / mesh_dimension_size))

    return num_entries