Computing useful input data structures to ease IP construction.

def _preprocess_input(self):
    """Computing useful input data structures to ease IP construction."""
    # Compute the sets of MTF dimensions used in operations/tensors.

    # a {string: frozenset(string)}, mapping operation name to MTF dimension
    # names.
    self._operation_name_to_mtf_dimension_set = {}
    # a {string: frozenset(string)}, mapping tensor name to MTF dimension names.
    self._tensor_name_to_mtf_dimension_set = {}

    for operation_name in self._graph.get_all_operation_names():
      self._operation_name_to_mtf_dimension_set[operation_name] = frozenset(
          set(self._graph.get_operation_mtf_dimension_names(
              operation_name)).intersection(
                  self._layout_validator.splittable_mtf_dimension_names))
    for tensor_name in self._graph.get_all_tensor_names():
      self._tensor_name_to_mtf_dimension_set[tensor_name] = frozenset(
          set(self._graph.get_tensor_mtf_dimension_names(tensor_name))
          .intersection(self._layout_validator.splittable_mtf_dimension_names))

    self._operation_mtf_dimension_sets = set(
        self._operation_name_to_mtf_dimension_set.values())
    self._mtf_dimension_sets = self._operation_mtf_dimension_sets | set(
        self._tensor_name_to_mtf_dimension_set.values())

    # Compute possible assignments for each set of MTF dimensions.
    self._assignments = {}  # indexed by MTF dimension set
    for mtf_dimension_set in self._mtf_dimension_sets:
      self._assignments[mtf_dimension_set] = _generate_assignments(
          mtf_dimension_set, self._layout_validator.mesh_dimension_name_to_size)