code: def get_tensor_num_entries(self, tensor_name, partial_layout=None,
                               mesh_dimension_to_size=None):
      
      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
docs: 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