import re


def _split_by_task(devices, values):
  """Partition devices and values by common task.

  Args:
    devices: list of device name strings
    values: list of T @{tf.tensor} of same length as devices.

  Returns:
    (per_task_devices, per_task_values) where both values are
    lists of lists with isomorphic structure: the outer list is
    indexed by task, and the inner list has length of the number
    of values belonging to that task.  per_task_devices contains
    the specific devices to which the values are local, and
    per_task_values contains the corresponding values.

  Raises:
    ValueError: devices must be same length as values.
  """
  num_devices = len(devices)
  if num_devices != len(values):
    raise ValueError("len(devices) must equal len(values)")
  pattern = re.compile(r"/task:(\d+)/")
  per_task_devices = []
  per_task_values = []
  for d in range(num_devices):
    m = pattern.search(devices[d])
    if m:
      index = int(m.group(1))
      while index >= len(per_task_devices):
        per_task_devices.append([])
        per_task_values.append([])
      per_task_devices[index].append(devices[d])
      per_task_values[index].append(values[d])
    else:
      assert False, "failed to parse device %s" % devices[d]
  return (per_task_devices, per_task_values)