import six
import itertools


def _generate_assignments(splittable_dimensions, mesh_dimension_to_size):
  """Generates all ways to map splittable dimensions to mesh dimensions.

  Args:
    splittable_dimensions: a frozenset of the names of splittable dimensions.
    mesh_dimension_to_size: a dictionary from mesh dimension name to size.

  Returns:
    A list of the valid assignments. Each assignment is a dict keyed by every
        splittable dimension, whose value is either a mesh dimension or None.
  """
  assignments = []
  for assignment_size in six.moves.xrange(
      1 + min(len(splittable_dimensions), len(mesh_dimension_to_size))):
    for s_dims_chosen in itertools.combinations(splittable_dimensions,
                                                assignment_size):
      for m_dims_chosen in itertools.permutations(mesh_dimension_to_size,
                                                  assignment_size):
        assignments.append(dict(zip(s_dims_chosen, m_dims_chosen)))
  return assignments