Tensor round to a number of decimal places

[wangjiawen2013]

Hi,
We can round a float in python:

import numpy as np
arr = np.array([1.23456789, 2.3456789, 3.456789])
arr_rounded = np.round(arr, 2)  # keep two decimal, output: [1.23, 2.35, 3.46]

how to round a tensor in burn like that in python ?

[laggui]

We have tensor.round() but the number of decimal places to round to is not configurable.

/edit: if you want to achieve that in the meantime, you can easily do it (though with additional steps)

pub fn round<B: Backend, const D: usize>(tensor: Tensor<B, D>, decimals: u32) -> Tensor<B, D> {
    if decimals == 0 {
        tensor.round()
    } else {
        let shift = 10u32.pow(decimals);
        tensor.mul_scalar(shift).round().div_scalar(shift)
    }
}

let x = Tensor::<B, 1>::from([1.23456789, 2.3456789, 3.456789]);
rounded = round(x, 2)

[wangjiawen2013]

It works! Why not adding this function to tensor.round() ?

[laggui]

This could be added to the tensor.round() as a default implementation, but as you can see it's not so difficult to achieve without 🙂