tenferro_runtime/
tensor.rs

1//! Concrete tensor operation extension trait.
2//!
3//! `tenferro-tensor` owns storage and backend traits. This runtime crate
4//! provides backend-parametric operation methods through [`TensorOpsExt`].
5
6use tenferro_ops::broadcast::{broadcast_input_plan, broadcast_shape, broadcast_shapes};
7use tenferro_tensor::validate::matmul_config_for_shapes;
8use tenferro_tensor::{CompareDir, DType, Error, Result, TensorBackend};
9
10use crate::TensorOpsExt;
11use tenferro_tensor::Tensor;
12
13impl TensorOpsExt for Tensor {
14    fn convert<B: TensorBackend>(&self, to: DType, backend: &mut B) -> Result<Tensor> {
15        convert(self, to, backend)
16    }
17
18    fn cast<B: TensorBackend>(&self, to: DType, backend: &mut B) -> Result<Tensor> {
19        cast(self, to, backend)
20    }
21
22    fn add<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
23        add(self, rhs, backend)
24    }
25
26    fn sub<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
27        sub(self, rhs, backend)
28    }
29
30    fn mul<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
31        mul(self, rhs, backend)
32    }
33
34    fn div<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
35        div(self, rhs, backend)
36    }
37
38    fn rem<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
39        rem(self, rhs, backend)
40    }
41
42    fn pow<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
43        pow(self, rhs, backend)
44    }
45
46    fn maximum<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
47        maximum(self, rhs, backend)
48    }
49
50    fn minimum<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
51        minimum(self, rhs, backend)
52    }
53
54    fn neg<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
55        neg(self, backend)
56    }
57
58    fn abs<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
59        abs(self, backend)
60    }
61
62    fn sign<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
63        sign(self, backend)
64    }
65
66    fn conj<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
67        conj(self, backend)
68    }
69
70    fn exp<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
71        exp(self, backend)
72    }
73
74    fn log<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
75        log(self, backend)
76    }
77
78    fn sin<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
79        sin(self, backend)
80    }
81
82    fn cos<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
83        cos(self, backend)
84    }
85
86    fn tanh<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
87        tanh(self, backend)
88    }
89
90    fn sqrt<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
91        sqrt(self, backend)
92    }
93
94    fn rsqrt<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
95        rsqrt(self, backend)
96    }
97
98    fn expm1<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
99        expm1(self, backend)
100    }
101
102    fn log1p<B: TensorBackend>(&self, backend: &mut B) -> Result<Tensor> {
103        log1p(self, backend)
104    }
105
106    fn compare<B: TensorBackend>(
107        &self,
108        rhs: &Tensor,
109        dir: CompareDir,
110        backend: &mut B,
111    ) -> Result<Tensor> {
112        compare(self, rhs, dir, backend)
113    }
114
115    fn where_select<B: TensorBackend>(
116        &self,
117        on_true: &Tensor,
118        on_false: &Tensor,
119        backend: &mut B,
120    ) -> Result<Tensor> {
121        where_select(self, on_true, on_false, backend)
122    }
123
124    fn clamp<B: TensorBackend>(
125        &self,
126        lower: &Tensor,
127        upper: &Tensor,
128        backend: &mut B,
129    ) -> Result<Tensor> {
130        clamp(self, lower, upper, backend)
131    }
132
133    fn matmul<B: TensorBackend>(&self, rhs: &Tensor, backend: &mut B) -> Result<Tensor> {
134        matmul(self, rhs, backend)
135    }
136
137    fn reshape<B: TensorBackend>(&self, shape: &[usize], backend: &mut B) -> Result<Tensor> {
138        reshape(self, shape, backend)
139    }
140
141    fn transpose<B: TensorBackend>(&self, perm: &[usize], backend: &mut B) -> Result<Tensor> {
142        transpose(self, perm, backend)
143    }
144
145    fn reduce_sum<B: TensorBackend>(&self, axes: &[usize], backend: &mut B) -> Result<Tensor> {
146        reduce_sum(self, axes, backend)
147    }
148}
149
150/// Convert a tensor to a different dtype using the checked conversion lattice.
151///
152/// Use [`cast`] for explicit lossy dtype projection.
153///
154/// # Examples
155///
156/// ```rust
157/// # use tenferro_cpu::CpuBackend;
158/// use tenferro_runtime::{DType, Tensor, TensorOpsExt};
159/// # let mut backend = CpuBackend::new();
160/// # let x = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
161/// let y = x.convert(DType::C64, &mut backend).unwrap();
162/// assert_eq!(y.dtype(), DType::C64);
163/// ```
164///
165/// # Errors
166///
167/// Returns an error when the requested conversion is outside tenferro's checked
168/// dtype-promotion lattice, or when the backend does not support the requested
169/// conversion.
170fn convert(input: &Tensor, to: DType, backend: &mut impl TensorBackend) -> Result<Tensor> {
171    backend.with_backend_session(|exec| exec.convert(input, to))
172}
173
174/// Cast a tensor to a different dtype using explicit dtype projection.
175///
176/// Unlike [`convert`], `cast` may truncate, narrow precision, project complex
177/// values to their real component, or use boolean truthiness where the backend
178/// supports the requested projection.
179///
180/// # Examples
181///
182/// ```rust
183/// # use tenferro_cpu::CpuBackend;
184/// use tenferro_runtime::{DType, Tensor, TensorOpsExt};
185/// # let mut backend = CpuBackend::new();
186/// # let x = Tensor::from_vec_col_major(vec![2], vec![1.2_f64, -2.8]).unwrap();
187/// let y = x.cast(DType::I32, &mut backend).unwrap();
188/// assert_eq!(y.as_slice::<i32>().unwrap(), &[1, -2]);
189/// ```
190///
191/// # Errors
192///
193/// Returns an error when the backend does not support the requested explicit
194/// dtype projection.
195fn cast(input: &Tensor, to: DType, backend: &mut impl TensorBackend) -> Result<Tensor> {
196    backend.with_backend_session(|exec| exec.cast(input, to))
197}
198
199/// Elementwise addition with NumPy-style broadcasting.
200///
201/// # Examples
202///
203/// ```rust
204/// # use tenferro_cpu::CpuBackend;
205/// use tenferro_runtime::{Tensor, TensorOpsExt};
206/// # let mut backend = CpuBackend::new();
207/// # let x = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
208/// # let y = Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap();
209/// let z = x.add(&y, &mut backend).unwrap();
210/// ```
211fn add(lhs: &Tensor, rhs: &Tensor, backend: &mut impl TensorBackend) -> Result<Tensor> {
212    let (lhs, rhs) = broadcast_binary(lhs, rhs, backend)?;
213    backend.with_backend_session(|exec| exec.add(&lhs, &rhs))
214}
215
216macro_rules! unary_fn {
217    ($name:ident, $method:ident, $summary:literal) => {
218        #[doc = $summary]
219        ///
220        /// # Examples
221        ///
222        /// ```rust
223        /// # use tenferro_cpu::CpuBackend;
224        /// use tenferro_runtime::{Tensor, TensorOpsExt};
225        /// # let mut backend = CpuBackend::new();
226        /// # let x = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 4.0]).unwrap();
227        #[doc = concat!("let y = x.", stringify!($name), "(&mut backend).unwrap();")]
228        /// ```
229        fn $name(input: &Tensor, backend: &mut impl TensorBackend) -> Result<Tensor> {
230            backend.with_backend_session(|exec| exec.$method(input))
231        }
232    };
233}
234
235macro_rules! binary_fn {
236    ($name:ident, $method:ident, $summary:literal) => {
237        #[doc = $summary]
238        ///
239        /// # Examples
240        ///
241        /// ```rust
242        /// # use tenferro_cpu::CpuBackend;
243        /// use tenferro_runtime::{Tensor, TensorOpsExt};
244        /// # let mut backend = CpuBackend::new();
245        /// # let x = Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 4.0]).unwrap();
246        /// # let y = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 8.0]).unwrap();
247        #[doc = concat!("let z = x.", stringify!($name), "(&y, &mut backend).unwrap();")]
248        /// ```
249        fn $name(lhs: &Tensor, rhs: &Tensor, backend: &mut impl TensorBackend) -> Result<Tensor> {
250            let (lhs, rhs) = broadcast_binary(lhs, rhs, backend)?;
251            backend.with_backend_session(|exec| exec.$method(&lhs, &rhs))
252        }
253    };
254}
255
256binary_fn!(
257    mul,
258    mul,
259    "Elementwise multiplication with NumPy-style broadcasting."
260);
261binary_fn!(
262    div,
263    div,
264    "Elementwise division with NumPy-style broadcasting."
265);
266binary_fn!(
267    rem,
268    rem,
269    "Elementwise remainder with NumPy-style broadcasting."
270);
271binary_fn!(pow, pow, "Elementwise power with NumPy-style broadcasting.");
272binary_fn!(
273    maximum,
274    maximum,
275    "Elementwise maximum with NumPy-style broadcasting."
276);
277binary_fn!(
278    minimum,
279    minimum,
280    "Elementwise minimum with NumPy-style broadcasting."
281);
282
283unary_fn!(neg, neg, "Elementwise negation.");
284unary_fn!(abs, abs, "Elementwise absolute value.");
285unary_fn!(sign, sign, "Elementwise sign.");
286unary_fn!(conj, conj, "Elementwise complex conjugate.");
287unary_fn!(exp, exp, "Elementwise exponential.");
288unary_fn!(log, log, "Elementwise natural logarithm.");
289unary_fn!(sin, sin, "Elementwise sine.");
290unary_fn!(cos, cos, "Elementwise cosine.");
291unary_fn!(tanh, tanh, "Elementwise hyperbolic tangent.");
292unary_fn!(sqrt, sqrt, "Elementwise square root.");
293unary_fn!(rsqrt, rsqrt, "Elementwise reciprocal square root.");
294unary_fn!(expm1, expm1, "Elementwise `exp(x) - 1`.");
295unary_fn!(log1p, log1p, "Elementwise `log(1 + x)`.");
296
297/// Elementwise subtraction with NumPy-style broadcasting.
298///
299/// # Examples
300///
301/// ```rust
302/// # use tenferro_cpu::CpuBackend;
303/// use tenferro_runtime::{Tensor, TensorOpsExt};
304/// # let mut backend = CpuBackend::new();
305/// # let x = Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 4.0]).unwrap();
306/// # let y = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 8.0]).unwrap();
307/// let z = x.sub(&y, &mut backend).unwrap();
308/// ```
309fn sub(lhs: &Tensor, rhs: &Tensor, backend: &mut impl TensorBackend) -> Result<Tensor> {
310    let (lhs, rhs) = broadcast_binary(lhs, rhs, backend)?;
311    backend.with_backend_session(|exec| exec.sub(&lhs, &rhs))
312}
313
314/// Elementwise comparison with NumPy-style broadcasting.
315///
316/// The result is a bool tensor.
317///
318/// # Examples
319///
320/// ```rust
321/// # use tenferro_cpu::CpuBackend;
322/// use tenferro_runtime::{CompareDir, Tensor, TensorOpsExt};
323/// # let mut backend = CpuBackend::new();
324/// # let x = Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 4.0]).unwrap();
325/// # let y = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 8.0]).unwrap();
326/// let z = x.compare(&y, CompareDir::Gt, &mut backend).unwrap();
327/// assert_eq!(z.as_slice::<bool>().unwrap(), &[true, false]);
328/// ```
329fn compare(
330    lhs: &Tensor,
331    rhs: &Tensor,
332    dir: CompareDir,
333    backend: &mut impl TensorBackend,
334) -> Result<Tensor> {
335    let (lhs, rhs) = broadcast_binary(lhs, rhs, backend)?;
336    backend.with_backend_session(|exec| exec.compare(&lhs, &rhs, &dir))
337}
338
339/// Select values from `on_true` or `on_false` using a condition tensor.
340///
341/// This corresponds to NumPy `where(condition, x, y)`.
342///
343/// # Examples
344///
345/// ```rust
346/// # use tenferro_cpu::CpuBackend;
347/// use tenferro_runtime::{CompareDir, Tensor, TensorOpsExt};
348/// # let mut backend = CpuBackend::new();
349/// # let x = Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 4.0]).unwrap();
350/// # let y = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 8.0]).unwrap();
351/// # let condition = x.compare(&y, CompareDir::Gt, &mut backend).unwrap();
352/// let z = condition.where_select(&x, &y, &mut backend).unwrap();
353/// ```
354fn where_select(
355    condition: &Tensor,
356    on_true: &Tensor,
357    on_false: &Tensor,
358    backend: &mut impl TensorBackend,
359) -> Result<Tensor> {
360    let (condition, on_true, on_false) = broadcast_ternary(condition, on_true, on_false, backend)?;
361    backend.with_backend_session(|exec| exec.select(&condition, &on_true, &on_false))
362}
363
364/// Clamp values elementwise between lower and upper bounds.
365///
366/// # Examples
367///
368/// ```rust
369/// # use tenferro_cpu::CpuBackend;
370/// use tenferro_runtime::{Tensor, TensorOpsExt};
371/// # let mut backend = CpuBackend::new();
372/// # let x = Tensor::from_vec_col_major(vec![2], vec![-2.0_f64, 4.0]).unwrap();
373/// # let lower = Tensor::from_vec_col_major(vec![], vec![0.0_f64]).unwrap();
374/// # let upper = Tensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
375/// let z = x.clamp(&lower, &upper, &mut backend).unwrap();
376/// ```
377fn clamp(
378    input: &Tensor,
379    lower: &Tensor,
380    upper: &Tensor,
381    backend: &mut impl TensorBackend,
382) -> Result<Tensor> {
383    let (input, lower, upper) = broadcast_ternary(input, lower, upper, backend)?;
384    backend.with_backend_session(|exec| exec.clamp(&input, &lower, &upper))
385}
386
387/// Matrix multiplication helper for rank-2 tensors.
388///
389/// This contracts the last dimension of `a` with the first dimension of `b`.
390///
391/// # Examples
392///
393/// ```rust
394/// # use tenferro_cpu::CpuBackend;
395/// use tenferro_runtime::{Tensor, TensorOpsExt};
396/// # let mut backend = CpuBackend::new();
397/// # let a = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
398/// # let b = Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64; 6]).unwrap();
399/// let c = a.matmul(&b, &mut backend).unwrap();
400/// ```
401fn matmul(a: &Tensor, b: &Tensor, backend: &mut impl TensorBackend) -> Result<Tensor> {
402    let config = matmul_config_for_shapes("matmul", a.shape(), b.shape())?;
403    backend.with_backend_session(|exec| exec.dot_general(a, b, &config))
404}
405
406/// Reshape a tensor without changing element order.
407///
408/// # Examples
409///
410/// ```rust
411/// # use tenferro_cpu::CpuBackend;
412/// use tenferro_runtime::{Tensor, TensorOpsExt};
413/// # let mut backend = CpuBackend::new();
414/// # let x = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
415/// let y = x.reshape(&[4], &mut backend).unwrap();
416/// assert_eq!(y.shape(), &[4]);
417/// ```
418fn reshape(input: &Tensor, shape: &[usize], backend: &mut impl TensorBackend) -> Result<Tensor> {
419    backend.with_backend_session(|exec| exec.reshape(input, shape))
420}
421
422/// Permute tensor axes.
423///
424/// # Examples
425///
426/// ```rust
427/// # use tenferro_cpu::CpuBackend;
428/// use tenferro_runtime::{Tensor, TensorOpsExt};
429/// # let mut backend = CpuBackend::new();
430/// # let x = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
431/// let y = x.transpose(&[1, 0], &mut backend).unwrap();
432/// assert_eq!(y.shape(), &[3, 2]);
433/// ```
434fn transpose(input: &Tensor, perm: &[usize], backend: &mut impl TensorBackend) -> Result<Tensor> {
435    backend.with_backend_session(|exec| exec.transpose(input, perm))
436}
437
438/// Sum a tensor over one or more axes.
439///
440/// # Examples
441///
442/// ```rust
443/// # use tenferro_cpu::CpuBackend;
444/// use tenferro_runtime::{Tensor, TensorOpsExt};
445/// # let mut backend = CpuBackend::new();
446/// # let x = Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
447/// let y = x.reduce_sum(&[0], &mut backend).unwrap();
448/// assert_eq!(y.shape(), &[2]);
449/// ```
450fn reduce_sum(input: &Tensor, axes: &[usize], backend: &mut impl TensorBackend) -> Result<Tensor> {
451    backend.with_backend_session(|exec| exec.reduce_sum(input, axes))
452}
453
454fn broadcast_binary(
455    lhs: &Tensor,
456    rhs: &Tensor,
457    backend: &mut impl TensorBackend,
458) -> Result<(Tensor, Tensor)> {
459    let shape = broadcast_shape(lhs.shape(), rhs.shape()).map_err(broadcast_error)?;
460    Ok((
461        broadcast_to(lhs, &shape, backend)?,
462        broadcast_to(rhs, &shape, backend)?,
463    ))
464}
465
466fn broadcast_ternary(
467    first: &Tensor,
468    second: &Tensor,
469    third: &Tensor,
470    backend: &mut impl TensorBackend,
471) -> Result<(Tensor, Tensor, Tensor)> {
472    let shape = broadcast_shapes([first.shape(), second.shape(), third.shape()])
473        .map_err(broadcast_error)?;
474    Ok((
475        broadcast_to(first, &shape, backend)?,
476        broadcast_to(second, &shape, backend)?,
477        broadcast_to(third, &shape, backend)?,
478    ))
479}
480
481fn broadcast_to(
482    input: &Tensor,
483    target_shape: &[usize],
484    backend: &mut impl TensorBackend,
485) -> Result<Tensor> {
486    let input_shape = input.shape();
487    if input_shape == target_shape {
488        return Ok(input.clone());
489    }
490
491    let plan = broadcast_input_plan(input_shape, target_shape).map_err(broadcast_error)?;
492    let source = if plan.source_shape == input_shape {
493        input.clone()
494    } else {
495        backend.with_backend_session(|exec| exec.reshape(input, &plan.source_shape))?
496    };
497    backend.with_backend_session(|exec| exec.broadcast_in_dim(&source, target_shape, &plan.dims))
498}
499
500fn broadcast_error(err: impl std::fmt::Display) -> Error {
501    Error::backend_failure("broadcast", err.to_string())
502}
tenferro_runtime/tensor.rs

tenferro_runtime/
tensor.rs
