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

tenferro_runtime/
tensor.rs
