Multi-dimensional step interpolation.

  Returns the multi-dimensional step interpolant to a function with
  given discrete data points (xp, fp), evaluated at x.

  Note that *N and *M indicate zero or more dimensions.

  Args:
    x: An array of shape [*N], the x-coordinates of the interpolated values.
    xp: An np.array of shape [D], the x-coordinates of the data points, must be
      increasing.
    fp: An np.array of shape [D, *M], the y-coordinates of the data points.
    **kwargs: Unused.

  Returns:
    An array of shape [*N, *M], the interpolated values.

def step_interpolation(x, xp, fp, **kwargs):
  """Multi-dimensional step interpolation.

  Returns the multi-dimensional step interpolant to a function with
  given discrete data points (xp, fp), evaluated at x.

  Note that *N and *M indicate zero or more dimensions.

  Args:
    x: An array of shape [*N], the x-coordinates of the interpolated values.
    xp: An np.array of shape [D], the x-coordinates of the data points, must be
      increasing.
    fp: An np.array of shape [D, *M], the y-coordinates of the data points.
    **kwargs: Unused.

  Returns:
    An array of shape [*N, *M], the interpolated values.
  """
  del kwargs  # Unused.
  xp = np.expand_dims(xp, -1)
  lower, upper = xp[:-1], xp[1:]
  conditions = (x >= lower) & (x < upper)
  # Underflow and overflow conditions and values. Values default to fp[0] and
  # fp[-1] respectively.
  conditions = np.concatenate([[x < xp[0]], conditions, [x >= xp[-1]]])
  values = np.concatenate([[fp[0]], fp])
  assert np.all(np.sum(conditions, 0) == 1), 'xp must be increasing.'
  indices = np.argmax(conditions, 0)
  return values[indices].astype(np.float32)