Collapse `shape` outside the interval (`a`,`b`).

  This function collapses `shape` outside the interval (`a`,`b`) by
  multiplying the dimensions before `a` into a single dimension,
  and mutliplying the dimensions after `b` into a single dimension.

  Args:
    shape: a tensor shape
    a: integer, position in shape
    b: integer, position in shape

  Returns:
    The collapsed shape, represented as a list.

  Examples:
    [1, 2, 3, 4, 5], (a=0, b=2) => [1, 1, 2, 60]
    [1, 2, 3, 4, 5], (a=1, b=3) => [1, 2, 3, 20]
    [1, 2, 3, 4, 5], (a=2, b=4) => [2, 3, 4, 5 ]
    [1, 2, 3, 4, 5], (a=3, b=5) => [6, 4, 5, 1 ]

def collapse_shape(shape, a, b):
  """Collapse `shape` outside the interval (`a`,`b`).

  This function collapses `shape` outside the interval (`a`,`b`) by
  multiplying the dimensions before `a` into a single dimension,
  and mutliplying the dimensions after `b` into a single dimension.

  Args:
    shape: a tensor shape
    a: integer, position in shape
    b: integer, position in shape

  Returns:
    The collapsed shape, represented as a list.

  Examples:
    [1, 2, 3, 4, 5], (a=0, b=2) => [1, 1, 2, 60]
    [1, 2, 3, 4, 5], (a=1, b=3) => [1, 2, 3, 20]
    [1, 2, 3, 4, 5], (a=2, b=4) => [2, 3, 4, 5 ]
    [1, 2, 3, 4, 5], (a=3, b=5) => [6, 4, 5, 1 ]
  """
  shape = list(shape)
  if a < 0:
    n_pad = -a
    pad = n_pad * [1]
    return collapse_shape(pad + shape, a + n_pad, b + n_pad)
  if b > len(shape):
    n_pad = b - len(shape)
    pad = n_pad * [1]
    return collapse_shape(shape + pad, a, b)
  return [product(shape[:a])] + shape[a:b] + [product(shape[b:])]