code: def step_interpolation(x, xp, fp, **kwargs):
    
    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)
docs: 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.