ruprim_host/binary/
dispatch_float.rs

1use ruda_core::tensor::{host::HostTensor, element::Scalar, TensorMetadata};
2use num_traits::ToPrimitive;
3use super::{BinaryOp, binary_op, scalar_op};
4#[cfg(not(feature = "std"))]
5#[allow(unused_imports)]
6use num_traits::Float;
7
8pub fn float_add(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
9    binary_op(lhs, rhs, |a, b| a + b, |a, b| a + b, Some(BinaryOp::Add))
10}
11
12pub fn float_add_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
13    let rhs_val = rhs.to_f64().unwrap();
14    scalar_op(lhs, rhs_val, |a, b| a + b, |a, b| a + b)
15}
16
17pub fn float_sub(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
18    binary_op(lhs, rhs, |a, b| a - b, |a, b| a - b, Some(BinaryOp::Sub))
19}
20
21pub fn float_sub_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
22    let rhs_val = rhs.to_f64().unwrap();
23    scalar_op(lhs, rhs_val, |a, b| a - b, |a, b| a - b)
24}
25
26pub fn float_mul(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
27    binary_op(lhs, rhs, |a, b| a * b, |a, b| a * b, Some(BinaryOp::Mul))
28}
29
30pub fn float_mul_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
31    let rhs_val = rhs.to_f64().unwrap();
32    scalar_op(lhs, rhs_val, |a, b| a * b, |a, b| a * b)
33}
34
35pub fn float_div(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
36    binary_op(lhs, rhs, |a, b| a / b, |a, b| a / b, Some(BinaryOp::Div))
37}
38
39pub fn float_div_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
40    let rhs_val = rhs.to_f64().unwrap();
41    scalar_op(lhs, rhs_val, |a, b| a / b, |a, b| a / b)
42}
43
44pub fn float_remainder(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
45    // Python/PyTorch-style remainder: result has same sign as divisor
46    binary_op(
47        lhs,
48        rhs,
49        |a, b| ((a % b) + b) % b,
50        |a, b| ((a % b) + b) % b,
51        None,
52    )
53}
54
55pub fn float_remainder_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
56    let rhs_val = rhs.to_f64().unwrap();
57    // Python/PyTorch-style remainder: result has same sign as divisor
58    scalar_op(
59        lhs,
60        rhs_val,
61        |a, b| ((a % b) + b) % b,
62        |a, b| ((a % b) + b) % b,
63    )
64}
65
66pub fn float_powf(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
67    binary_op(lhs, rhs, |a: f32, b| a.powf(b), |a: f64, b| a.powf(b), None)
68}
69
70pub fn float_powf_scalar_impl(tensor: HostTensor, value: Scalar) -> HostTensor {
71    let exp = value.to_f64().unwrap();
72    scalar_op(tensor, exp, |a: f32, b| a.powf(b), |a: f64, b| a.powf(b))
73}
74
75pub fn float_atan2(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
76    binary_op(
77        lhs,
78        rhs,
79        |a: f32, b| a.atan2(b),
80        |a: f64, b| a.atan2(b),
81        None,
82    )
83}
84
85pub fn float_hypot(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
86    binary_op(lhs, rhs, |a: f32, b| a.hypot(b), |a: f64, b| a.hypot(b), None)
87}
88
89pub fn float_powi(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
90    super::integer_power::tensor(lhs, rhs)
91}
92
93pub fn float_powi_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
94    if let Scalar::UInt(exponent) = rhs
95        && exponent > i64::MAX as u64
96    {
97        return super::integer_power::scalar(lhs, false, exponent);
98    }
99    match rhs.to_i64().unwrap() {
100        0 => crate::fill::float_ones(lhs.shape(), lhs.dtype().into()),
101        1 => lhs,
102        2 => crate::binary::dispatch_float::float_mul(lhs.clone(), lhs),
103        -1 => crate::unary::recip(lhs),
104        -2 => crate::unary::recip(crate::binary::dispatch_float::float_mul(lhs.clone(), lhs)),
105        exponent => super::integer_power::scalar(lhs, exponent < 0, exponent.unsigned_abs()),
106    }
107}
108
109pub fn float_powf_scalar(tensor: HostTensor, value: Scalar) -> HostTensor {
110    if let Some(exp) = value.try_as_integer() {
111        crate::binary::dispatch_float::float_powi_scalar(tensor, exp)
112    } else {
113        crate::binary::dispatch_float::float_powf_scalar_impl(tensor, value)
114    }
115}
116
ruprim_host/binary/dispatch_float.rs

ruprim_host/binary/
dispatch_float.rs
