Function searchsorted

utils/stepfun.py:6–25

Find indices where v should be inserted into a to maintain order. Args: a: tensor, the sorted reference points that we are scanning to see where v should lie. v: tensor, the query points that we are pretending to insert into a. Does not need to be sorted. All but the last dime

(a, v)

def searchsorted(a, v):
    """Find indices where v should be inserted into a to maintain order.

  Args:
    a: tensor, the sorted reference points that we are scanning to see where v
      should lie.
    v: tensor, the query points that we are pretending to insert into a. Does
      not need to be sorted. All but the last dimensions should match or expand
      to those of a, the last dimension can differ.

  Returns:
    (idx_lo, idx_hi), where a[idx_lo] <= v < a[idx_hi], unless v is out of the
    range [a[0], a[-1]] in which case idx_lo and idx_hi are both the first or
    last index of a.
  """
    i = torch.arange(a.shape[-1], device=a.device)
    v_ge_a = v[..., None, :] >= a[..., :, None]
    idx_lo = torch.max(torch.where(v_ge_a, i[..., :, None], i[..., :1, None]), -2).values
    idx_hi = torch.min(torch.where(~v_ge_a, i[..., :, None], i[..., -1:, None]), -2).values
    return idx_lo, idx_hi


def query(tq, t, y, outside_value=0):