nerfacc.searchsorted

nerfacc.searchsorted(sorted_sequence, values)

Searchsorted that supports flattened tensor.

This function returns {ids_left, ids_right} such that:

sorted_sequence.vals.gather(-1, ids_left) <= values.vals < sorted_sequence.vals.gather(-1, ids_right)

Note

When values is out of range of sorted_sequence, we return the corresponding ids as if the values is clipped to the range of sorted_sequence. See the example below.

Parameters:

• sorted_sequence (Union[RayIntervals, RaySamples]) – A RayIntervals or RaySamples object. We assume the sorted_sequence.vals is acendingly sorted for each ray.

• values (Union[RayIntervals, RaySamples]) – A RayIntervals or RaySamples object.

Returns:

• ids_left: A LongTensor with the same shape as values.vals.

• ids_right: A LongTensor with the same shape as values.vals.

Return type:

A tuple of LongTensor

Example

>>> sorted_sequence = RayIntervals(
...     vals=torch.tensor([0.0, 1.0, 0.0, 1.0, 2.0], device="cuda"),
...     packed_info=torch.tensor([[0, 2], [2, 3]], device="cuda"),
... )
>>> values = RayIntervals(
...     vals=torch.tensor([0.5, 1.5, 2.5], device="cuda"),
...     packed_info=torch.tensor([[0, 1], [1, 2]], device="cuda"),
... )
>>> ids_left, ids_right = searchsorted(sorted_sequence, values)
>>> ids_left
tensor([0, 3, 3], device='cuda:0')
>>> ids_right
tensor([1, 4, 4], device='cuda:0')
>>> sorted_sequence.vals.gather(-1, ids_left)
tensor([0., 1., 1.], device='cuda:0')
>>> sorted_sequence.vals.gather(-1, ids_right)
tensor([1., 2., 2.], device='cuda:0')