Skip to main content

tenferro_runtime/
typed_tensor.rs

1//! Typed tensor operation extension traits.
2//!
3//! Operation families that are no longer part of core, including einsum, live
4//! in their extension crates.
5
6use tenferro_ops::broadcast::{
7    broadcast_input_plan, broadcast_shape, broadcast_shapes, BroadcastError,
8};
9use tenferro_tensor::validate::matmul_config_for_shapes;
10use tenferro_tensor::{
11    CompareDir, DType, Error, Result, Tensor, TensorBackend, TensorRead, TensorScalar,
12    ValidationError,
13};
14
15use crate::{TypedTensorMaskOpsExt, TypedTensorOpsExt};
16use tenferro_tensor::TypedTensor;
17
18impl<T: TensorScalar> TypedTensorOpsExt<T> for TypedTensor<T> {
19    fn add<B: TensorBackend>(
20        &self,
21        rhs: &TypedTensor<T>,
22        backend: &mut B,
23    ) -> Result<TypedTensor<T>> {
24        add(self, rhs, backend)
25    }
26
27    fn sub<B: TensorBackend>(
28        &self,
29        rhs: &TypedTensor<T>,
30        backend: &mut B,
31    ) -> Result<TypedTensor<T>> {
32        sub(self, rhs, backend)
33    }
34
35    fn mul<B: TensorBackend>(
36        &self,
37        rhs: &TypedTensor<T>,
38        backend: &mut B,
39    ) -> Result<TypedTensor<T>> {
40        mul(self, rhs, backend)
41    }
42
43    fn div<B: TensorBackend>(
44        &self,
45        rhs: &TypedTensor<T>,
46        backend: &mut B,
47    ) -> Result<TypedTensor<T>> {
48        div(self, rhs, backend)
49    }
50
51    fn rem<B: TensorBackend>(
52        &self,
53        rhs: &TypedTensor<T>,
54        backend: &mut B,
55    ) -> Result<TypedTensor<T>> {
56        rem(self, rhs, backend)
57    }
58
59    fn pow<B: TensorBackend>(
60        &self,
61        rhs: &TypedTensor<T>,
62        backend: &mut B,
63    ) -> Result<TypedTensor<T>> {
64        pow(self, rhs, backend)
65    }
66
67    fn maximum<B: TensorBackend>(
68        &self,
69        rhs: &TypedTensor<T>,
70        backend: &mut B,
71    ) -> Result<TypedTensor<T>> {
72        maximum(self, rhs, backend)
73    }
74
75    fn minimum<B: TensorBackend>(
76        &self,
77        rhs: &TypedTensor<T>,
78        backend: &mut B,
79    ) -> Result<TypedTensor<T>> {
80        minimum(self, rhs, backend)
81    }
82
83    fn neg<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
84        neg(self, backend)
85    }
86
87    fn abs<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
88        abs(self, backend)
89    }
90
91    fn sign<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
92        sign(self, backend)
93    }
94
95    fn conj<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
96        conj(self, backend)
97    }
98
99    fn exp<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
100        exp(self, backend)
101    }
102
103    fn log<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
104        log(self, backend)
105    }
106
107    fn sin<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
108        sin(self, backend)
109    }
110
111    fn cos<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
112        cos(self, backend)
113    }
114
115    fn tanh<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
116        tanh(self, backend)
117    }
118
119    fn sqrt<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
120        sqrt(self, backend)
121    }
122
123    fn rsqrt<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
124        rsqrt(self, backend)
125    }
126
127    fn expm1<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
128        expm1(self, backend)
129    }
130
131    fn log1p<B: TensorBackend>(&self, backend: &mut B) -> Result<TypedTensor<T>> {
132        log1p(self, backend)
133    }
134
135    fn compare<B: TensorBackend>(
136        &self,
137        rhs: &TypedTensor<T>,
138        dir: CompareDir,
139        backend: &mut B,
140    ) -> Result<TypedTensor<bool>> {
141        compare(self, rhs, dir, backend)
142    }
143
144    fn clamp<B: TensorBackend>(
145        &self,
146        lower: &TypedTensor<T>,
147        upper: &TypedTensor<T>,
148        backend: &mut B,
149    ) -> Result<TypedTensor<T>> {
150        clamp(self, lower, upper, backend)
151    }
152
153    fn matmul<B: TensorBackend>(
154        &self,
155        rhs: &TypedTensor<T>,
156        backend: &mut B,
157    ) -> Result<TypedTensor<T>> {
158        matmul(self, rhs, backend)
159    }
160
161    fn reduce_sum<B: TensorBackend>(
162        &self,
163        axes: &[usize],
164        backend: &mut B,
165    ) -> Result<TypedTensor<T>> {
166        reduce_sum(self, axes, backend)
167    }
168
169    fn reshape<B: TensorBackend>(
170        &self,
171        shape: &[usize],
172        backend: &mut B,
173    ) -> Result<TypedTensor<T>> {
174        reshape(self, shape, backend)
175    }
176
177    fn transpose<B: TensorBackend>(
178        &self,
179        perm: &[usize],
180        backend: &mut B,
181    ) -> Result<TypedTensor<T>> {
182        transpose(self, perm, backend)
183    }
184
185    fn broadcast_in_dim<B: TensorBackend>(
186        &self,
187        shape: &[usize],
188        dims: &[usize],
189        backend: &mut B,
190    ) -> Result<TypedTensor<T>> {
191        broadcast_in_dim(self, shape, dims, backend)
192    }
193}
194
195impl TypedTensorMaskOpsExt for TypedTensor<bool> {
196    fn where_select<T: TensorScalar, B: TensorBackend>(
197        &self,
198        on_true: &TypedTensor<T>,
199        on_false: &TypedTensor<T>,
200        backend: &mut B,
201    ) -> Result<TypedTensor<T>> {
202        where_select(self, on_true, on_false, backend)
203    }
204}
205
206/// Elementwise addition with NumPy-style broadcasting.
207///
208/// # Examples
209///
210/// ```rust
211/// # use tenferro_cpu::CpuBackend;
212/// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
213/// # let mut backend = CpuBackend::new();
214/// # let x = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 2.0]).unwrap();
215/// # let y = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![3.0, 4.0]).unwrap();
216/// let z = x.add(&y, &mut backend).unwrap();
217/// ```
218fn add<T: TensorScalar>(
219    lhs: &TypedTensor<T>,
220    rhs: &TypedTensor<T>,
221    backend: &mut impl TensorBackend,
222) -> Result<TypedTensor<T>> {
223    let (lhs, rhs) = broadcast_binary_read(lhs, rhs, backend)?;
224    let out =
225        backend.with_backend_session(|exec| exec.add_read(lhs.tensor_read(), rhs.tensor_read()))?;
226    into_typed_result("add", out)
227}
228
229macro_rules! unary_fn {
230    ($name:ident, $method:ident, $summary:literal) => {
231        #[doc = $summary]
232        ///
233        /// # Examples
234        ///
235        /// ```rust
236        /// # use tenferro_cpu::CpuBackend;
237        /// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
238        /// # let mut backend = CpuBackend::new();
239        /// # let x = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 4.0]).unwrap();
240        #[doc = concat!("let y = x.", stringify!($name), "(&mut backend).unwrap();")]
241        /// ```
242        fn $name<T: TensorScalar>(
243            input: &TypedTensor<T>,
244            backend: &mut impl TensorBackend,
245        ) -> Result<TypedTensor<T>> {
246            let out = backend.with_backend_session(|exec| exec.$method(T::tensor_read(input)))?;
247            into_typed_result(stringify!($name), out)
248        }
249    };
250}
251
252macro_rules! binary_fn {
253    ($name:ident, $method:ident, $summary:literal) => {
254        #[doc = $summary]
255        ///
256        /// # Examples
257        ///
258        /// ```rust
259        /// # use tenferro_cpu::CpuBackend;
260        /// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
261        /// # let mut backend = CpuBackend::new();
262        /// # let x = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 4.0]).unwrap();
263        /// # let y = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 8.0]).unwrap();
264        #[doc = concat!("let z = x.", stringify!($name), "(&y, &mut backend).unwrap();")]
265        /// ```
266        fn $name<T: TensorScalar>(
267            lhs: &TypedTensor<T>,
268            rhs: &TypedTensor<T>,
269            backend: &mut impl TensorBackend,
270        ) -> Result<TypedTensor<T>> {
271            let (lhs, rhs) = broadcast_binary_read(lhs, rhs, backend)?;
272            let out = backend
273                .with_backend_session(|exec| exec.$method(lhs.tensor_read(), rhs.tensor_read()))?;
274            into_typed_result(stringify!($name), out)
275        }
276    };
277}
278
279binary_fn!(
280    mul,
281    mul_read,
282    "Elementwise multiplication with NumPy-style broadcasting."
283);
284binary_fn!(
285    div,
286    div_read,
287    "Elementwise division with NumPy-style broadcasting."
288);
289binary_fn!(
290    rem,
291    rem_read,
292    "Elementwise remainder with NumPy-style broadcasting."
293);
294binary_fn!(
295    pow,
296    pow_read,
297    "Elementwise power with NumPy-style broadcasting."
298);
299binary_fn!(
300    maximum,
301    maximum_read,
302    "Elementwise maximum with NumPy-style broadcasting."
303);
304binary_fn!(
305    minimum,
306    minimum_read,
307    "Elementwise minimum with NumPy-style broadcasting."
308);
309
310unary_fn!(neg, neg_read, "Elementwise negation.");
311unary_fn!(abs, abs_read, "Elementwise absolute value.");
312unary_fn!(sign, sign_read, "Elementwise sign.");
313unary_fn!(conj, conj_read, "Elementwise complex conjugate.");
314unary_fn!(exp, exp_read, "Elementwise exponential.");
315unary_fn!(log, log_read, "Elementwise natural logarithm.");
316unary_fn!(sin, sin_read, "Elementwise sine.");
317unary_fn!(cos, cos_read, "Elementwise cosine.");
318unary_fn!(tanh, tanh_read, "Elementwise hyperbolic tangent.");
319unary_fn!(sqrt, sqrt_read, "Elementwise square root.");
320unary_fn!(rsqrt, rsqrt_read, "Elementwise reciprocal square root.");
321unary_fn!(expm1, expm1_read, "Elementwise `exp(x) - 1`.");
322unary_fn!(log1p, log1p_read, "Elementwise `log(1 + x)`.");
323
324/// Elementwise subtraction with NumPy-style broadcasting.
325///
326/// # Examples
327///
328/// ```rust
329/// # use tenferro_cpu::CpuBackend;
330/// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
331/// # let mut backend = CpuBackend::new();
332/// # let x = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 4.0]).unwrap();
333/// # let y = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 8.0]).unwrap();
334/// let z = x.sub(&y, &mut backend).unwrap();
335/// ```
336fn sub<T: TensorScalar>(
337    lhs: &TypedTensor<T>,
338    rhs: &TypedTensor<T>,
339    backend: &mut impl TensorBackend,
340) -> Result<TypedTensor<T>> {
341    let (lhs, rhs) = broadcast_binary_read(lhs, rhs, backend)?;
342    let out =
343        backend.with_backend_session(|exec| exec.sub_read(lhs.tensor_read(), rhs.tensor_read()))?;
344    into_typed_result("sub", out)
345}
346
347/// Elementwise comparison with NumPy-style broadcasting.
348///
349/// The result is a bool tensor.
350///
351/// # Examples
352///
353/// ```rust
354/// # use tenferro_cpu::CpuBackend;
355/// use tenferro_runtime::{CompareDir, TypedTensor, TypedTensorMaskOpsExt, TypedTensorOpsExt};
356/// # let mut backend = CpuBackend::new();
357/// # let x = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 4.0]).unwrap();
358/// # let y = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 8.0]).unwrap();
359/// let z = x.compare(&y, CompareDir::Gt, &mut backend).unwrap();
360/// assert_eq!(z.host_data().unwrap(), &[true, false]);
361/// ```
362fn compare<T: TensorScalar>(
363    lhs: &TypedTensor<T>,
364    rhs: &TypedTensor<T>,
365    dir: CompareDir,
366    backend: &mut impl TensorBackend,
367) -> Result<TypedTensor<bool>> {
368    let (lhs, rhs) = broadcast_binary_read(lhs, rhs, backend)?;
369    let out = backend.with_backend_session(|exec| {
370        exec.compare_read(lhs.tensor_read(), rhs.tensor_read(), &dir)
371    })?;
372    into_typed_result("compare", out)
373}
374
375/// Select values from `on_true` or `on_false` using a condition tensor.
376///
377/// This corresponds to NumPy `where(condition, x, y)`.
378///
379/// # Examples
380///
381/// ```rust
382/// # use tenferro_cpu::CpuBackend;
383/// use tenferro_runtime::{CompareDir, TypedTensor, TypedTensorMaskOpsExt, TypedTensorOpsExt};
384/// # let mut backend = CpuBackend::new();
385/// # let x = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![2.0, 4.0]).unwrap();
386/// # let y = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![1.0, 8.0]).unwrap();
387/// # let condition = x.compare(&y, CompareDir::Gt, &mut backend).unwrap();
388/// let z = condition.where_select(&x, &y, &mut backend).unwrap();
389/// ```
390fn where_select<T: TensorScalar>(
391    condition: &TypedTensor<bool>,
392    on_true: &TypedTensor<T>,
393    on_false: &TypedTensor<T>,
394    backend: &mut impl TensorBackend,
395) -> Result<TypedTensor<T>> {
396    let (condition, on_true, on_false) =
397        broadcast_ternary_read(condition, on_true, on_false, backend)?;
398    let out = backend.with_backend_session(|exec| {
399        exec.select_read(
400            condition.tensor_read(),
401            on_true.tensor_read(),
402            on_false.tensor_read(),
403        )
404    })?;
405    into_typed_result("where_select", out)
406}
407
408/// Clamp values elementwise between lower and upper bounds.
409///
410/// # Examples
411///
412/// ```rust
413/// # use tenferro_cpu::CpuBackend;
414/// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
415/// # let mut backend = CpuBackend::new();
416/// # let x = TypedTensor::<f64>::from_vec_col_major(vec![2], vec![-2.0, 4.0]).unwrap();
417/// # let lower = TypedTensor::<f64>::from_vec_col_major(vec![], vec![0.0]).unwrap();
418/// # let upper = TypedTensor::<f64>::from_vec_col_major(vec![], vec![3.0]).unwrap();
419/// let z = x.clamp(&lower, &upper, &mut backend).unwrap();
420/// ```
421fn clamp<T: TensorScalar>(
422    input: &TypedTensor<T>,
423    lower: &TypedTensor<T>,
424    upper: &TypedTensor<T>,
425    backend: &mut impl TensorBackend,
426) -> Result<TypedTensor<T>> {
427    let (input, lower, upper) = broadcast_ternary_read(input, lower, upper, backend)?;
428    let out = backend.with_backend_session(|exec| {
429        exec.clamp_read(
430            input.tensor_read(),
431            lower.tensor_read(),
432            upper.tensor_read(),
433        )
434    })?;
435    into_typed_result("clamp", out)
436}
437
438/// Matrix multiplication helper for rank-2 typed tensors.
439///
440/// This contracts the last dimension of `a` with the first dimension of `b`.
441///
442/// # Examples
443///
444/// ```rust
445/// # use tenferro_cpu::CpuBackend;
446/// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
447/// # let mut backend = CpuBackend::new();
448/// # let a = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
449/// # let b = TypedTensor::<f64>::from_vec_col_major(vec![3, 2], vec![1.0; 6]).unwrap();
450/// let c = a.matmul(&b, &mut backend).unwrap();
451/// ```
452fn matmul<T: TensorScalar>(
453    a: &TypedTensor<T>,
454    b: &TypedTensor<T>,
455    backend: &mut impl TensorBackend,
456) -> Result<TypedTensor<T>> {
457    let config = matmul_config_for_shapes("matmul", a.shape(), b.shape())?;
458    let out = backend.with_backend_session(|exec| {
459        exec.dot_general_read(T::tensor_read(a), T::tensor_read(b), &config)
460    })?;
461    into_typed_result("matmul", out)
462}
463
464/// Sum elements across one or more axes.
465///
466/// # Examples
467///
468/// ```rust
469/// # use tenferro_cpu::CpuBackend;
470/// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
471/// # let mut backend = CpuBackend::new();
472/// let x = TypedTensor::<f64>::from_vec_col_major(
473///     vec![2, 3],
474///     vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0],
475/// )?;
476/// let row_sums = x.reduce_sum(&[1], &mut backend).unwrap();
477/// assert_eq!(row_sums.host_data()?, &[6.0, 15.0]);
478/// # Ok::<(), tenferro_runtime::Error>(())
479/// ```
480fn reduce_sum<T: TensorScalar>(
481    input: &TypedTensor<T>,
482    axes: &[usize],
483    backend: &mut impl TensorBackend,
484) -> Result<TypedTensor<T>> {
485    let out =
486        backend.with_backend_session(|exec| exec.reduce_sum_read(T::tensor_read(input), axes))?;
487    into_typed_result("reduce_sum", out)
488}
489
490/// Reshape a typed tensor through the backend structural operation.
491///
492/// # Examples
493///
494/// ```rust
495/// # use tenferro_cpu::CpuBackend;
496/// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
497/// # let mut backend = CpuBackend::new();
498/// let x = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
499/// let y = x.reshape(&[3, 2], &mut backend).unwrap();
500/// assert_eq!(y.shape(), &[3, 2]);
501/// ```
502fn reshape<T: TensorScalar>(
503    input: &TypedTensor<T>,
504    shape: &[usize],
505    backend: &mut impl TensorBackend,
506) -> Result<TypedTensor<T>> {
507    let out =
508        backend.with_backend_session(|exec| exec.reshape_read(T::tensor_read(input), shape))?;
509    into_typed_result("reshape", out)
510}
511
512/// Permute typed tensor axes through the backend structural operation.
513///
514/// # Examples
515///
516/// ```rust
517/// # use tenferro_cpu::CpuBackend;
518/// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
519/// # let mut backend = CpuBackend::new();
520/// let x = TypedTensor::<f64>::from_vec_col_major(vec![2, 3], vec![1.0; 6]).unwrap();
521/// let y = x.transpose(&[1, 0], &mut backend).unwrap();
522/// assert_eq!(y.shape(), &[3, 2]);
523/// ```
524fn transpose<T: TensorScalar>(
525    input: &TypedTensor<T>,
526    perm: &[usize],
527    backend: &mut impl TensorBackend,
528) -> Result<TypedTensor<T>> {
529    let out =
530        backend.with_backend_session(|exec| exec.transpose_read(T::tensor_read(input), perm))?;
531    into_typed_result("transpose", out)
532}
533
534/// Broadcast a typed tensor into a larger shape.
535///
536/// `dims` maps each input axis to its output axis, following the concrete
537/// backend `broadcast_in_dim` contract.
538///
539/// # Examples
540///
541/// ```rust
542/// # use tenferro_cpu::CpuBackend;
543/// use tenferro_runtime::{TypedTensor, TypedTensorOpsExt};
544/// # let mut backend = CpuBackend::new();
545/// let row = TypedTensor::<f64>::from_vec_col_major(vec![3], vec![1.0, 2.0, 3.0]).unwrap();
546/// let matrix = row.broadcast_in_dim(&[2, 3], &[1], &mut backend).unwrap();
547/// assert_eq!(matrix.shape(), &[2, 3]);
548/// ```
549fn broadcast_in_dim<T: TensorScalar>(
550    input: &TypedTensor<T>,
551    shape: &[usize],
552    dims: &[usize],
553    backend: &mut impl TensorBackend,
554) -> Result<TypedTensor<T>> {
555    let out = backend.with_backend_session(|exec| {
556        exec.broadcast_in_dim_read(T::tensor_read(input), shape, dims)
557    })?;
558    into_typed_result("broadcast_in_dim", out)
559}
560
561enum ReadInput<'a> {
562    Borrowed(TensorRead<'a>),
563    Owned(Tensor),
564}
565
566impl ReadInput<'_> {
567    fn tensor_read(&self) -> TensorRead<'_> {
568        match self {
569            Self::Borrowed(read) => read.clone(),
570            Self::Owned(tensor) => TensorRead::from_tensor(tensor),
571        }
572    }
573}
574
575fn broadcast_binary_read<'a, T: TensorScalar>(
576    lhs: &'a TypedTensor<T>,
577    rhs: &'a TypedTensor<T>,
578    backend: &mut impl TensorBackend,
579) -> Result<(ReadInput<'a>, ReadInput<'a>)> {
580    let shape = broadcast_shape(lhs.shape(), rhs.shape()).map_err(broadcast_error)?;
581    Ok((
582        broadcast_to_read(lhs, &shape, backend)?,
583        broadcast_to_read(rhs, &shape, backend)?,
584    ))
585}
586
587fn broadcast_ternary_read<'a, C: TensorScalar, T: TensorScalar>(
588    first: &'a TypedTensor<C>,
589    second: &'a TypedTensor<T>,
590    third: &'a TypedTensor<T>,
591    backend: &mut impl TensorBackend,
592) -> Result<(ReadInput<'a>, ReadInput<'a>, ReadInput<'a>)> {
593    let shape = broadcast_shapes([first.shape(), second.shape(), third.shape()])
594        .map_err(broadcast_error)?;
595    Ok((
596        broadcast_to_read(first, &shape, backend)?,
597        broadcast_to_read(second, &shape, backend)?,
598        broadcast_to_read(third, &shape, backend)?,
599    ))
600}
601
602fn broadcast_to_read<'a, T: TensorScalar>(
603    input: &'a TypedTensor<T>,
604    target_shape: &[usize],
605    backend: &mut impl TensorBackend,
606) -> Result<ReadInput<'a>> {
607    if input.shape() == target_shape {
608        return Ok(ReadInput::Borrowed(T::tensor_read(input)));
609    }
610
611    let plan = broadcast_input_plan(input.shape(), target_shape).map_err(broadcast_error)?;
612    let source = if plan.source_shape == input.shape() {
613        ReadInput::Borrowed(T::tensor_read(input))
614    } else {
615        let reshaped = backend.with_backend_session(|exec| {
616            exec.reshape_read(T::tensor_read(input), &plan.source_shape)
617        })?;
618        ReadInput::Owned(reshaped)
619    };
620    let out = backend.with_backend_session(|exec| {
621        exec.broadcast_in_dim_read(source.tensor_read(), target_shape, &plan.dims)
622    })?;
623    Ok(ReadInput::Owned(out))
624}
625
626fn broadcast_error(err: BroadcastError) -> Error {
627    match err {
628        BroadcastError::IncompatibleBinary { lhs, rhs } => {
629            Error::shape_mismatch("broadcast", lhs, rhs)
630        }
631        BroadcastError::IncompatibleInput { input, output } => {
632            Error::shape_mismatch("broadcast", input, output)
633        }
634        BroadcastError::RankTooLarge { input, output } => {
635            Error::rank_mismatch("broadcast", output.len(), input.len())
636        }
637    }
638}
639
640fn into_typed_result<T: TensorScalar>(op: &'static str, tensor: Tensor) -> Result<TypedTensor<T>> {
641    let actual = tensor.dtype();
642    T::into_typed(tensor).map_err(|_| {
643        Error::validation(
644            op,
645            ValidationError::DTypeMismatch {
646                expected: core_dtype(T::dtype()),
647                actual: core_dtype(actual),
648            },
649        )
650    })
651}
652
653fn core_dtype(dtype: DType) -> tenferro_tensor::core::DType {
654    match dtype {
655        DType::F32 => tenferro_tensor::core::DType::F32,
656        DType::F64 => tenferro_tensor::core::DType::F64,
657        DType::I32 => tenferro_tensor::core::DType::I32,
658        DType::I64 => tenferro_tensor::core::DType::I64,
659        DType::Bool => tenferro_tensor::core::DType::Bool,
660        DType::C32 => tenferro_tensor::core::DType::C32,
661        DType::C64 => tenferro_tensor::core::DType::C64,
662    }
663}