code: def get_tensor_size(self, tensor_name, partial_layout=None,
                        mesh_dimension_to_size=None):
      
      return (self.get_tensor_dtype(tensor_name).size *
              self.get_tensor_num_entries(tensor_name, partial_layout,
                                          mesh_dimension_to_size))
docs: The size of a tensor in bytes.
  
      If partial_layout is specified, then mesh_dimension_to_size must also be. In
      this case, the size 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