Impl round to even for candle

crates/burn-candle/src/ops/tensor.rs

@@ -338,7 +338,25 @@ impl<F: FloatCandleElement, I: IntCandleElement> FloatTensorOps<Self> for Candle
     }
 
     fn float_round(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
-        CandleTensor::new(tensor.tensor.round().unwrap())
+        let inner = |tensor: FloatTensor<Self>| -> candle_core::Result<FloatTensor<Self>> {
+            // implements round_to_even for consistent behavior vs libtorch
+            // https://github.com/pytorch/pytorch/blob/main/torch/csrc/jit/runtime/register_ops_utils.h#L65-L67
+
+            let floor_a = tensor.tensor.floor()?;
+            let frac_part = tensor.tensor.sub(&floor_a)?;
+
+            let half = (candle_core::Tensor::ones_like(&tensor.tensor)? * 0.5)?;
+            let mask_half = frac_part.eq(&half)?;
+            let half_tensor = tensor.tensor.mul(&half)?;
+            let rounded_half = half_tensor.round()?;
+            let doubled =
+                rounded_half.mul(&(candle_core::Tensor::ones_like(&tensor.tensor)? * 2.0)?)?;
+            let standard_round = tensor.tensor.round()?;
+            Ok(CandleTensor::new(
+                mask_half.where_cond(&doubled, &standard_round)?,
+            ))
+        };
+        inner(tensor).unwrap()
     }
 
     fn float_floor(tensor: FloatTensor<Self>) -> FloatTensor<Self> {