code: def evaluate_layout(self, layout):
      
      layout_dict = {}
      if layout:
        for pair in layout.split(";"):
          mtf_dimension_name, mesh_dimension_name = pair.split(":", 1)
          if (mtf_dimension_name in
              self._layout_validator.splittable_mtf_dimension_names):
            layout_dict[mtf_dimension_name] = mesh_dimension_name
          else:
            logging.warning("Skipping unsplittable dimension %s.",
                            mtf_dimension_name)
  
      tensor_memory = {}  # {string: float}, size of each tensor under our layout
      for tensor_name in self._graph.get_all_tensor_names():
        if self._graph.is_tensor_on_canonical_device(tensor_name):
          tensor_memory[tensor_name] = self._graph.get_tensor_size(
              tensor_name, layout_dict,
              self._layout_validator.mesh_dimension_name_to_size)
        else:
          tensor_memory[tensor_name] = 0.0
  
      peak_memory_usage = 0.0
      for tensor_names in self._get_memory_contents():
        memory_usage = 0.0
        for tensor_name in tensor_names:
          memory_usage += tensor_memory[tensor_name]
        peak_memory_usage = max(peak_memory_usage, memory_usage)
      return peak_memory_usage
docs: The current objective value for the given layout.
  
      TODO(joshuawang): The current function does not check that the given
      layout is valid.
  
      Args:
        layout: a string, representing a layout to evaluate (e.g.
            "d_ff:m1;heads:m2").
  
      Returns:
        A float, the objective value.