kopitiam_tensor/tensor/
elementwise.rs

1//! Broadcasting elementwise arithmetic: `add`, `sub`, `mul`, `div`.
2//!
3//! All four share one implementation — the only thing that differs between
4//! them is which `f32, f32 -> f32` function to apply — so they are thin
5//! wrappers around [`broadcast_binary`] rather than four copies of the same
6//! broadcast-then-zip-then-collect logic.
7
8use kopitiam_core::{DType, Result};
9
10use crate::storage::Storage;
11
12use super::Tensor;
13
14impl Tensor {
15    /// Elementwise `self + other`, broadcasting shapes per
16    /// [`kopitiam_core::Shape::broadcast`].
17    pub fn add(&self, other: &Tensor) -> Result<Tensor> {
18        broadcast_binary(self, other, |a, b| a + b)
19    }
20
21    pub fn sub(&self, other: &Tensor) -> Result<Tensor> {
22        broadcast_binary(self, other, |a, b| a - b)
23    }
24
25    pub fn mul(&self, other: &Tensor) -> Result<Tensor> {
26        broadcast_binary(self, other, |a, b| a * b)
27    }
28
29    /// Elementwise `self / other`. Division by zero follows ordinary IEEE
30    /// 754 float semantics (`+/-inf` or `NaN`) rather than erroring — a
31    /// tensor op has no way to know whether that is a bug or an expected
32    /// edge case (e.g. a masked-out attention score), so it is left to the
33    /// caller to detect if it matters.
34    pub fn div(&self, other: &Tensor) -> Result<Tensor> {
35        broadcast_binary(self, other, |a, b| a / b)
36    }
37}
38
39fn broadcast_binary(a: &Tensor, b: &Tensor, f: impl Fn(f32, f32) -> f32) -> Result<Tensor> {
40    a.require_dtype(DType::F32)?;
41    b.require_dtype(DType::F32)?;
42    let out_shape = a.shape.broadcast(&b.shape)?;
43    let a_view = a.broadcast_to(out_shape.clone())?;
44    let b_view = b.broadcast_to(out_shape.clone())?;
45    let Storage::F32(a_data) = a_view.storage.as_ref() else { unreachable!() };
46    let Storage::F32(b_data) = b_view.storage.as_ref() else { unreachable!() };
47    let result: Vec<f32> = a_view
48        .logical_offsets()
49        .zip(b_view.logical_offsets())
50        .map(|(ia, ib)| f(a_data[ia], b_data[ib]))
51        .collect();
52    Tensor::from_f32(result, out_shape)
53}
54
55#[cfg(test)]
56mod tests {
57    use kopitiam_core::Error;
58
59    use super::*;
60
61    #[test]
62    fn add_same_shape_matches_hand_computation() {
63        let a = Tensor::from_f32(vec![1.0, 2.0, 3.0], [3]).unwrap();
64        let b = Tensor::from_f32(vec![10.0, 20.0, 30.0], [3]).unwrap();
65        assert_eq!(a.add(&b).unwrap().to_vec_f32().unwrap(), vec![11.0, 22.0, 33.0]);
66    }
67
68    #[test]
69    fn mul_broadcasts_a_row_vector_over_a_matrix() {
70        // [[1,2,3],[4,5,6]] * [10,20,30] (broadcast over rows)
71        let a = Tensor::from_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]).unwrap();
72        let b = Tensor::from_f32(vec![10.0, 20.0, 30.0], [3]).unwrap();
73        let c = a.mul(&b).unwrap();
74        assert_eq!(c.shape().dims(), &[2, 3]);
75        assert_eq!(c.to_vec_f32().unwrap(), vec![10.0, 40.0, 90.0, 40.0, 100.0, 180.0]);
76    }
77
78    #[test]
79    fn sub_broadcasts_a_column_vector_over_a_matrix() {
80        // [[1,2],[3,4]] - [[10],[20]] -> [[-9,-8],[-17,-16]]
81        let a = Tensor::from_f32(vec![1.0, 2.0, 3.0, 4.0], [2, 2]).unwrap();
82        let b = Tensor::from_f32(vec![10.0, 20.0], [2, 1]).unwrap();
83        let c = a.sub(&b).unwrap();
84        assert_eq!(c.to_vec_f32().unwrap(), vec![-9.0, -8.0, -17.0, -16.0]);
85    }
86
87    #[test]
88    fn div_by_a_scalar_broadcasts_against_rank_zero() {
89        let a = Tensor::from_f32(vec![2.0, 4.0, 6.0], [3]).unwrap();
90        let scalar = Tensor::from_f32(vec![2.0], []).unwrap();
91        assert_eq!(a.div(&scalar).unwrap().to_vec_f32().unwrap(), vec![1.0, 2.0, 3.0]);
92    }
93
94    #[test]
95    fn broadcasting_incompatible_shapes_is_rejected() {
96        let a = Tensor::from_f32(vec![1.0, 2.0, 3.0], [3]).unwrap();
97        let b = Tensor::from_f32(vec![1.0, 2.0], [2]).unwrap();
98        assert!(matches!(a.add(&b), Err(Error::NotBroadcastable { .. })));
99    }
100
101    #[test]
102    fn elementwise_ops_reject_non_f32_operands() {
103        let a = Tensor::from_i32(vec![1, 2, 3], [3]).unwrap();
104        let b = Tensor::from_i32(vec![1, 2, 3], [3]).unwrap();
105        assert!(matches!(a.add(&b), Err(Error::DTypeMismatch { .. })));
106    }
107}
kopitiam_tensor/tensor/elementwise.rs

kopitiam_tensor/tensor/
elementwise.rs
