Function _ring_2d

tensorflow/python/tpu/device_assignment.py:178–212

Ring-order of a height x width mesh. For example, in a 4x4 mesh, this returns the following order. 0 -- 1 -- 2 -- 3 | | | | 15-- 6 -- 5 -- 4 | | | | 14-- 7 -- 8 -- 9 | | | | 13-- 12-- 11-- 10 Args: height: An integer represents the hei

(height, width)

def _ring_2d(height, width):
  """Ring-order of a height x width mesh.

  For example, in a 4x4 mesh, this returns the following order.
    0 -- 1 -- 2 -- 3
    |    |    |    |
    15-- 6 -- 5 -- 4
    |    |    |    |
    14-- 7 -- 8 -- 9
    |    |    |    |
    13-- 12-- 11-- 10

  Args:
    height: An integer represents the height.
    width: An integer represents the width.

  Returns:
    A list of [y, x] pairs with ring order.
  """
  if height == 1:
    return [(0, i) for i in range(width)]
  if width == 1:
    return [(i, 0) for i in range(height)]
  if height % 2 != 0:
    logging.warning("Odd dimension")
    return [(i % height, i // height) for i in range(width * height)]
  ret = [(0, 0)]
  for i in range(height // 2):
    for j in range(1, width):
      ret.append((2 * i, j))
    for j in range(width - 1, 0, -1):
      ret.append((2 * i + 1, j))
  for i in range(height - 1, 0, -1):
    ret.append((i, 0))
  return ret


def device_assignment(topology,