def _range_along_dimension(range_dim, shape):
  """Construct a Tensor whose values are the index along a dimension.

  Construct a Tensor that counts the distance along a single dimension. This is
  useful, for example, when constructing an identity matrix,

    >>> x = _range_along_dimension(0, [2, 2]).eval()
    >>> x
    array([[0, 0],
           [1, 1]], dtype=int32)

    >>> y = _range_along_dimension(1, [2, 2]).eval()
    >>> y
    array([[0, 1],
           [0, 1]], dtype=int32)

    >>> tf.cast(tf.equal(x, y), dtype=tf.int32).eval()
    array([[1, 0],
           [0, 1]], dtype=int32)

  Args:
    range_dim: int. Dimension to count indices on.
    shape: 1D Tensor of ints. Shape of Tensor to construct.

  Returns:
    A Tensor whose values are the same as the range along dimension range_dim.

  Raises:
    ValueError: If range_dim isn't a valid dimension.
  """
  rank = len(shape)
  if range_dim >= rank:
    raise ValueError("Cannot calculate range along non-existent index.")
  indices = tf.range(start=0, limit=shape[range_dim])
  indices = tf.reshape(
      indices,
      shape=[1 if i != range_dim else shape[range_dim] for i in range(rank)])
  return tf.tile(indices,
                 [shape[i] if i != range_dim else 1 for i in range(rank)])
Construct a Tensor whose values are the index along a dimension.

  Construct a Tensor that counts the distance along a single dimension. This is
  useful, for example, when constructing an identity matrix,

    >>> x = _range_along_dimension(0, [2, 2]).eval()
    >>> x
    array([[0, 0],
           [1, 1]], dtype=int32)

    >>> y = _range_along_dimension(1, [2, 2]).eval()
    >>> y
    array([[0, 1],
           [0, 1]], dtype=int32)

    >>> tf.cast(tf.equal(x, y), dtype=tf.int32).eval()
    array([[1, 0],
           [0, 1]], dtype=int32)

  Args:
    range_dim: int. Dimension to count indices on.
    shape: 1D Tensor of ints. Shape of Tensor to construct.

  Returns:
    A Tensor whose values are the same as the range along dimension range_dim.

  Raises:
    ValueError: If range_dim isn't a valid dimension.