def _split(v, col1, col2):
  """Split a vector into a vector + a vector or array.

  The first vector is 1D with col1 columns.  The second has col2 columns and
  is a 1-D vector if len(v) == col1 + col2 or an array otherwise.

  Args:
    v: vector to split
    col1: number of columns for the first portion
    col2: number of columns for the second portion

  Returns:
    A tuple consisting of the first split vector and the second.
  """
  v1 = v[:col1]
  v2 = v[col1:]
  if len(v2) == col2:
    return v1, v2
  return v1, v2.reshape([-1, col2])