What is the correct implementation for rounding a 2D tensor given the rounding values in a 1D tensor?

Question

This is what I have done so far:

def round_values(predictions):
    # Rounding values
    rounded = torch.tensor([1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0])

    # Find nearest rounding value 
    dif = predictions.unsqueeze(2) - rounded.view(1,1,-1)
    indexes = torch.argmin(abs(dif), dim=2)

    # Fill tensor with nearest round value
    # TO IMPROVE
    rounded_predictions = torch.zeros(*predictions.shape)
    for i in range(rounded_predictions.shape[0]):
        for j in range(rounded_predictions.shape[1]):
            index = indexes[i,j]
            rounded_predictions[i,j] = rounded[index]
        
    return rounded_predictions 

a = 4*torch.rand((6,6)) + 1
print(a)
print(round_values(a))

The double loop is terrible. I would like something like rounded[indexes], so it returns a tensor with the same shape as indexes with the values of rounded, but I have not found a way to do it at tensor level instead of element-wise. I would like Any idea?

Answer 1

A thousand times faster for a 100 x 100 array (on cpu)

def round_values_bob(predictions):
    return torch.clamp(torch.round(predictions * 2) / 2, 1, 5)