Initializing the variables of the IP.

def _initialize_variables(self):
    """Initializing the variables of the IP."""
    # Initialize global variables.
    self._global_vars = {}  # Indexed by (MTF dimension, mesh dimension)
    for mtf_dimension_name in (
        self._layout_validator.splittable_mtf_dimension_names):
      for mesh_dimension_name in (
          self._layout_validator.mesh_dimension_name_to_size):
        name = _global_var_name(mtf_dimension_name, mesh_dimension_name)
        self._global_vars[(mtf_dimension_name, mesh_dimension_name)] = (
            self._model.NewBoolVar(name))

    # Initialize local variables.
    self._local_vars = {}  # Indexed by (tensorflow dimension set), then name of
    # assignment.
    for mtf_dimension_set in self._mtf_dimension_sets:
      self._local_vars[mtf_dimension_set] = {}
      for assignment in self._assignments[mtf_dimension_set]:
        # TODO(joshuawang): Avoid hash collision no matter what dimension names
        # are; don't hash by this local var name, swap to using a tuple encoding
        # of the full assignment instead.
        name = _local_var_name(mtf_dimension_set, assignment)
        self._local_vars[mtf_dimension_set][name] = (
            self._model.NewBoolVar(name))

    # Initialize memory variable. We need a crude upper bound on memory, so we
    # use the total size of all tensors under the empty assignment.
    # NOTE(joshuawang): This bound could be improved by factoring in the
    # schedule.
    memory_upper_bound = 0
    for tensor_name in self._graph.get_all_tensor_names():
      if self._graph.is_tensor_on_canonical_device(tensor_name):
        memory_upper_bound += int(self._graph.get_tensor_size(tensor_name))
    self._memory_var = self._model.NewIntVar(0, memory_upper_bound, "z")