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

tenferro_runtime/
tensor.rs
