[Python] pytorch: torch.bucketize with not-strictly increasing sequence

I am working with a batched interpolation class where a typical task is to find the knot index for each evaluation point in the batch. This sounds like a job for torch.bucketize, but the knot vector is not strictly increasing but has repeated vals. See the example below.

The documentation of torch.bucketize states for the boundary argument: "...must contain a strictly increasing sequence, or the return value is undefined."

So is there a similar function that can be used instead? Or is the documentation not clear enough at this point and instead, only the returned indices on the left and right extreme can be incorrect?

import torch

if __name__ == "__main__":
# Repeated knots at start and end (typical for splines)
boundaries = torch.tensor([0.0, 0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0, 1.0])

# Evaluation points around boundaries and middle
x_test = torch.tensor([-0.1, 0.0, 0.25, 0.5, 0.75, 1.0, 1.1], dtype=torch.float64)

indices = torch.bucketize(x_test, boundaries, right=False)

print("Knots: ", knots.tolist())
print("x values: ", x_test.tolist())
print("Indices: ", indices.tolist())

Continue reading...