From 5d7af72ba00d4b5ed3571e1ff418f61198b3b898 Mon Sep 17 00:00:00 2001 From: Jani Monoses Date: Wed, 10 Jun 2026 08:30:28 +0300 Subject: [PATCH] Round operatio --- catgrad/src/category/core.rs | 2 + catgrad/src/category/lang/ops.rs | 4 ++ catgrad/src/interpreter/backend/candle.rs | 58 +++++++++++++++++++ catgrad/src/interpreter/backend/mod.rs | 1 + catgrad/src/interpreter/backend/ndarray.rs | 12 ++++ catgrad/src/interpreter/backend/shape_only.rs | 4 ++ catgrad/src/interpreter/tensor_op.rs | 1 + catgrad/src/pass/to_core.rs | 1 + catgrad/src/stdlib/ops.rs | 2 +- 9 files changed, 84 insertions(+), 1 deletion(-) diff --git a/catgrad/src/category/core.rs b/catgrad/src/category/core.rs index ce0bf2f6..a195e279 100644 --- a/catgrad/src/category/core.rs +++ b/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), } } diff --git a/catgrad/src/category/lang/ops.rs b/catgrad/src/category/lang/ops.rs index ac41c7fe..98786990 100644 --- a/catgrad/src/category/lang/ops.rs +++ b/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 diff --git a/catgrad/src/interpreter/backend/candle.rs b/catgrad/src/interpreter/backend/candle.rs index 8c74a7d1..4ca1d6dd 100644 --- a/catgrad/src/interpreter/backend/candle.rs +++ b/catgrad/src/interpreter/backend/candle.rs @@ -826,6 +826,17 @@ impl Backend for CandleBackend { } } + fn round(&self, x: TaggedTensor) -> TaggedTensor { + 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) -> TaggedTensor { 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 = 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::().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::() + .unwrap() + .iter() + .zip(expected.iter()) + .enumerate() + { + assert_eq!( + actual, expected, + "Mismatch at index {i}: got {actual}, expected {expected}" + ); + } +} diff --git a/catgrad/src/interpreter/backend/mod.rs b/catgrad/src/interpreter/backend/mod.rs index 2e7e69b4..12973b06 100644 --- a/catgrad/src/interpreter/backend/mod.rs +++ b/catgrad/src/interpreter/backend/mod.rs @@ -75,6 +75,7 @@ pub trait Backend: Clone + Debug { fn exp(&self, x: TaggedTensor) -> TaggedTensor; fn log(&self, x: TaggedTensor) -> TaggedTensor; fn floor(&self, x: TaggedTensor) -> TaggedTensor; + fn round(&self, x: TaggedTensor) -> TaggedTensor; fn neg(&self, x: TaggedTensor) -> TaggedTensor; fn broadcast(&self, x: TaggedTensor, shape: Shape) -> TaggedTensor; fn reshape(&self, x: TaggedTensor, new_shape: Shape) -> TaggedTensor; diff --git a/catgrad/src/interpreter/backend/ndarray.rs b/catgrad/src/interpreter/backend/ndarray.rs index e83a06ca..cca588e7 100644 --- a/catgrad/src/interpreter/backend/ndarray.rs +++ b/catgrad/src/interpreter/backend/ndarray.rs @@ -398,6 +398,14 @@ impl Backend for NdArrayBackend { } } + fn round(&self, x: TaggedTensor) -> TaggedTensor { + 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) -> TaggedTensor { use TaggedTensorTuple::*; match x { @@ -635,6 +643,10 @@ impl NdArrayBackend { .map_collect(|&a, &b| a.powf(b)) } + fn round_f32(x: ArrayD) -> ArrayD { + x.mapv(f32::round_ties_even) + } + fn pow_u32(x: ArrayD, y: ArrayD) -> ArrayD { ndarray::Zip::from(&x) .and(&y) diff --git a/catgrad/src/interpreter/backend/shape_only.rs b/catgrad/src/interpreter/backend/shape_only.rs index 97eb6e8a..aeef60e3 100644 --- a/catgrad/src/interpreter/backend/shape_only.rs +++ b/catgrad/src/interpreter/backend/shape_only.rs @@ -201,6 +201,10 @@ impl Backend for ShapeOnlyBackend { x } + fn round(&self, x: TaggedTensor) -> TaggedTensor { + x + } + fn neg(&self, x: TaggedTensor) -> TaggedTensor { x } diff --git a/catgrad/src/interpreter/tensor_op.rs b/catgrad/src/interpreter/tensor_op.rs index 97b06d0a..30817b8c 100644 --- a/catgrad/src/interpreter/tensor_op.rs +++ b/catgrad/src/interpreter/tensor_op.rs @@ -21,6 +21,7 @@ pub(crate) fn tensor_op( 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), diff --git a/catgrad/src/pass/to_core.rs b/catgrad/src/pass/to_core.rs index 5a4e5136..424858c5 100644 --- a/catgrad/src/pass/to_core.rs +++ b/catgrad/src/pass/to_core.rs @@ -105,6 +105,7 @@ pub(crate) fn core_declarations() -> HashMap { (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))), diff --git a/catgrad/src/stdlib/ops.rs b/catgrad/src/stdlib/ops.rs index 415b6f7b..2297c5f8 100644 --- a/catgrad/src/stdlib/ops.rs +++ b/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 {