Construct a linear trajectory from x0.

  Args:
    x0: N-D float tensor.
    velocity: N-D float tensor
    t: [sequence_length]-length float tensor

  Returns:
    x: [sequence_length, ndims] float tensor.

def _get_linear_trajectory(x0, velocity, t):
  """Construct a linear trajectory from x0.

  Args:
    x0: N-D float tensor.
    velocity: N-D float tensor
    t: [sequence_length]-length float tensor

  Returns:
    x: [sequence_length, ndims] float tensor.
  """
  x0 = tf.convert_to_tensor(x0)
  velocity = tf.convert_to_tensor(velocity)
  t = tf.convert_to_tensor(t)
  if x0.shape.ndims != 1:
    raise ValueError("x0 must be a rank 1 tensor")
  if velocity.shape.ndims != 1:
    raise ValueError("velocity must be a rank 1 tensor")
  if t.shape.ndims != 1:
    raise ValueError("t must be a rank 1 tensor")
  x0 = tf.expand_dims(x0, axis=0)
  velocity = tf.expand_dims(velocity, axis=0)
  dx = velocity * tf.expand_dims(t, axis=-1)
  linear_trajectories = x0 + dx
  assert linear_trajectories.shape.ndims == 2, \
    "linear_trajectories should be a rank 2 tensor"
  return linear_trajectories