Round operation

catgrad/src/category/core.rs

@@ -259,6 +259,7 @@ pub enum ScalarOp {
     Exp,   // 1 → 1
     Log,   // 1 → 1
     Floor, // 1 → 1
+    Round, // 1 → 1
     Where, // 3 → 1
 }
 
@@ -282,6 +283,7 @@ impl ScalarOp {
             ScalarOp::Exp => (1, 1),
             ScalarOp::Log => (1, 1),
             ScalarOp::Floor => (1, 1),
+            ScalarOp::Round => (1, 1),
             ScalarOp::Where => (3, 1),
         }
     }

catgrad/src/category/lang/ops.rs

@@ -100,6 +100,10 @@ pub fn floor(builder: &Builder, value: Var) -> Var {
     var::fn_operation(builder, &[value], Object::Tensor, op!["tensor", "floor"])
 }
 
+pub fn round(builder: &Builder, value: Var) -> Var {
+    var::fn_operation(builder, &[value], Object::Tensor, op!["tensor", "round"])
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Declarations
 

catgrad/src/interpreter/backend/candle.rs

@@ -826,6 +826,17 @@ impl Backend for CandleBackend {
         }
     }
 
+    fn round(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
+        use TaggedTensorTuple::*;
+        match x {
+            F32([arr]) => F32([Self::unary_eager(arr, DType::F32, Self::round)]),
+            F16([arr]) => F16([Self::unary_eager(arr, DType::F16, Self::round)]),
+            BF16([arr]) => BF16([Self::unary_eager(arr, DType::BF16, Self::round)]),
+            FP8([arr]) => FP8([Self::unary_eager(arr, DType::F8E4M3, Self::round)]),
+            _ => panic!("Invalid type for round"),
+        }
+    }
+
     fn max(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
         use TaggedTensorTuple::*;
         match x {
@@ -1365,6 +1376,23 @@ impl CandleBackend {
         x.floor().unwrap().into()
     }
 
+    fn round(x: &Tensor) -> CandleTensor {
+        let dtype = x.dtype();
+        let shape = x.dims().to_vec();
+        let device = x.device().clone();
+        let result_vec: Vec<f32> = Self::float_tensor_to_f32_vec(x)
+            .into_iter()
+            .map(f32::round_ties_even)
+            .collect();
+        let result_tensor = Tensor::from_vec(result_vec, shape, &device).unwrap();
+        let result_tensor = if dtype == DType::F32 {
+            result_tensor
+        } else {
+            result_tensor.to_dtype(dtype).unwrap()
+        };
+        result_tensor.into()
+    }
+
     // Candle's pow function does not support negative base and silently generates NaNs
     // so we do element-wise powf https://github.com/huggingface/candle/issues/1640
     fn pow(x: &Tensor, y: &Tensor) -> CandleTensor {
@@ -1575,3 +1603,33 @@ fn test_indexed_select_rhs_matmul_matches_materialized_gather() {
         expected.flatten_all().unwrap().to_vec1::<f32>().unwrap()
     );
 }
+
+#[test]
+fn test_round_matches_python_ties_even() {
+    let tensor = Tensor::new(
+        &[1.2f32, 1.8, 2.5, 3.5, -1.2, -1.8, -2.5, -3.5],
+        &candle_core::Device::Cpu,
+    )
+    .unwrap()
+    .reshape(&[2, 4])
+    .unwrap();
+
+    let actual = CandleBackend::round(&tensor).materialize();
+    let expected = [1.0f32, 2.0, 2.0, 4.0, -1.0, -2.0, -2.0, -4.0];
+
+    assert_eq!(actual.dims(), &[2, 4]);
+    for (i, (&actual, &expected)) in actual
+        .flatten_all()
+        .unwrap()
+        .to_vec1::<f32>()
+        .unwrap()
+        .iter()
+        .zip(expected.iter())
+        .enumerate()
+    {
+        assert_eq!(
+            actual, expected,
+            "Mismatch at index {i}: got {actual}, expected {expected}"
+        );
+    }
+}

catgrad/src/interpreter/backend/mod.rs

@@ -75,6 +75,7 @@ pub trait Backend: Clone + Debug {
     fn exp(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
     fn log(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
     fn floor(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
+    fn round(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
     fn neg(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self>;
     fn broadcast(&self, x: TaggedTensor<Self>, shape: Shape) -> TaggedTensor<Self>;
     fn reshape(&self, x: TaggedTensor<Self>, new_shape: Shape) -> TaggedTensor<Self>;

catgrad/src/interpreter/backend/ndarray.rs

@@ -398,6 +398,14 @@ impl Backend for NdArrayBackend {
         }
     }
 
+    fn round(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
+        use TaggedTensorTuple::*;
+        match x {
+            F32([arr]) => from_f32(Self::round_f32(arr.unwrap_f32())),
+            _ => panic!("Invalid input types for round"),
+        }
+    }
+
     fn max(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
         use TaggedTensorTuple::*;
         match x {
@@ -635,6 +643,10 @@ impl NdArrayBackend {
             .map_collect(|&a, &b| a.powf(b))
     }
 
+    fn round_f32(x: ArrayD<f32>) -> ArrayD<f32> {
+        x.mapv(f32::round_ties_even)
+    }
+
     fn pow_u32(x: ArrayD<u32>, y: ArrayD<u32>) -> ArrayD<u32> {
         ndarray::Zip::from(&x)
             .and(&y)

catgrad/src/interpreter/backend/shape_only.rs

@@ -201,6 +201,10 @@ impl Backend for ShapeOnlyBackend {
         x
     }
 
+    fn round(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
+        x
+    }
+
     fn neg(&self, x: TaggedTensor<Self>) -> TaggedTensor<Self> {
         x
     }

catgrad/src/interpreter/tensor_op.rs

@@ -21,6 +21,7 @@ pub(crate) fn tensor_op<B: Backend>(
         TensorOp::Map(ScalarOp::Exp) => unary_op(backend, args, ssa, B::exp),
         TensorOp::Map(ScalarOp::Log) => unary_op(backend, args, ssa, B::log),
         TensorOp::Map(ScalarOp::Floor) => unary_op(backend, args, ssa, B::floor),
+        TensorOp::Map(ScalarOp::Round) => unary_op(backend, args, ssa, B::round),
         TensorOp::Map(ScalarOp::Neg) => unary_op(backend, args, ssa, B::neg),
         TensorOp::Map(ScalarOp::Mul) => binop(backend, args, ssa, B::mul),
         TensorOp::Map(ScalarOp::Div) => binop(backend, args, ssa, B::div),

catgrad/src/pass/to_core.rs

@@ -105,6 +105,7 @@ pub(crate) fn core_declarations() -> HashMap<lang::Path, core::Operation> {
         (path!["tensor", "exp"], Operation::Tensor(Map(Exp))),
         (path!["tensor", "log"], Operation::Tensor(Map(Log))),
         (path!["tensor", "floor"], Operation::Tensor(Map(Floor))),
+        (path!["tensor", "round"], Operation::Tensor(Map(Round))),
         (path!["tensor", "lt"], Operation::Tensor(Map(LT))),
         (path!["tensor", "gt"], Operation::Tensor(Map(GT))),
         (path!["tensor", "gte"], Operation::Tensor(Map(GTE))),

catgrad/src/stdlib/ops.rs

@@ -6,7 +6,7 @@ use crate::prelude::{Builder, Var};
 pub use ops::{
     arange, argmax, broadcast, cast, concat, cond, cos, dtype, dtype_constant, eq, exp, floor, gt,
     gte, index, log, lt, lte, matmul, max, nat, nat_to_u32, pack, param, pow, probe, reshape,
-    shape, sin, slice, sum, topk, transpose, unpack, where_cond,
+    round, shape, sin, slice, sum, topk, transpose, unpack, where_cond,
 };
 
 pub fn get(builder: &Builder, dim: impl IntoNatVar, start: impl IntoNatVar, x: Var) -> Var {