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 num_complex::Complex64;
7use tenferro_ops::broadcast::{broadcast_shape, broadcast_shapes, BroadcastError};
8use tenferro_tensor::validate::matmul_config_for_shapes;
9use tenferro_tensor::{
10    ActivationOp, CompareDir, DType, DotGeneralConfig, Engine, Error, Result, Tensor, TensorRead,
11    TensorScalar, TensorValue, ValidationError,
12};
13use tinyvec::TinyVec;
14
15use crate::composite;
16use crate::composite::session::{borrowed, run_session_composite, typed_borrowed};
17use crate::{TypedTensorMaskSessionOpsExt, TypedTensorSessionOpsExt};
18use tenferro_tensor::TypedTensor;
19
20impl<T: TensorScalar> TypedTensorSessionOpsExt<T> for TypedTensor<T> {
21    fn add(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
22        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
23        let out = session.add_read(lhs.tensor_read(), rhs.tensor_read())?;
24        into_typed_result("add", out)
25    }
26
27    fn add_into(
28        &self,
29        rhs: &TypedTensor<T>,
30        out: &mut TypedTensor<T>,
31        session: &mut dyn Engine,
32    ) -> Result<()> {
33        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
34        session.add_read_into(lhs.tensor_read(), rhs.tensor_read(), T::tensor_write(out))
35    }
36
37    fn mul(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
38        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
39        let out = session.mul_read(lhs.tensor_read(), rhs.tensor_read())?;
40        into_typed_result("mul", out)
41    }
42
43    fn exp(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
44        let out = session.exp_read(T::tensor_read(self))?;
45        into_typed_result("exp", out)
46    }
47
48    fn reduce_sum(
49        &self,
50        axes: Option<&[usize]>,
51        session: &mut dyn Engine,
52    ) -> Result<TypedTensor<T>> {
53        let axes = crate::tensor::all_axes_if_none(self.shape().len(), axes);
54        let out = session.reduce_sum_read(T::tensor_read(self), &axes)?;
55        into_typed_result("reduce_sum", out)
56    }
57
58    fn sub(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
59        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
60        let out = session.sub_read(lhs.tensor_read(), rhs.tensor_read())?;
61        into_typed_result("sub", out)
62    }
63
64    fn div(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
65        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
66        let out = session.div_read(lhs.tensor_read(), rhs.tensor_read())?;
67        into_typed_result("div", out)
68    }
69
70    fn rem(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
71        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
72        let out = session.rem_read(lhs.tensor_read(), rhs.tensor_read())?;
73        into_typed_result("rem", out)
74    }
75
76    fn pow(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
77        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
78        let out = session.pow_read(lhs.tensor_read(), rhs.tensor_read())?;
79        into_typed_result("pow", out)
80    }
81
82    fn maximum(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
83        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
84        let out = session.maximum_read(lhs.tensor_read(), rhs.tensor_read())?;
85        into_typed_result("maximum", out)
86    }
87
88    fn minimum(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
89        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
90        let out = session.minimum_read(lhs.tensor_read(), rhs.tensor_read())?;
91        into_typed_result("minimum", out)
92    }
93
94    fn neg(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
95        let out = session.neg_read(T::tensor_read(self))?;
96        into_typed_result("neg", out)
97    }
98
99    fn abs(&self, session: &mut dyn Engine) -> Result<TypedTensor<T::Real>> {
100        let out = session.abs_read(T::tensor_read(self))?;
101        into_typed_result::<T::Real>("abs", out)
102    }
103
104    fn sign(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
105        let out = session.sign_read(T::tensor_read(self))?;
106        into_typed_result("sign", out)
107    }
108
109    fn conj(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
110        let out = session.conj_read(T::tensor_read(self))?;
111        into_typed_result("conj", out)
112    }
113
114    fn log(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
115        let out = session.log_read(T::tensor_read(self))?;
116        into_typed_result("log", out)
117    }
118
119    fn expm1(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
120        let out = session.expm1_read(T::tensor_read(self))?;
121        into_typed_result("expm1", out)
122    }
123
124    fn log1p(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
125        let out = session.log1p_read(T::tensor_read(self))?;
126        into_typed_result("log1p", out)
127    }
128
129    fn erf(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
130        let out = session.erf_read(T::tensor_read(self))?;
131        into_typed_result("erf", out)
132    }
133
134    fn sin(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
135        let out = session.sin_read(T::tensor_read(self))?;
136        into_typed_result("sin", out)
137    }
138
139    fn cos(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
140        let out = session.cos_read(T::tensor_read(self))?;
141        into_typed_result("cos", out)
142    }
143
144    fn tanh(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
145        let out = session.tanh_read(T::tensor_read(self))?;
146        into_typed_result("tanh", out)
147    }
148
149    fn sqrt(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
150        let out = session.sqrt_read(T::tensor_read(self))?;
151        into_typed_result("sqrt", out)
152    }
153
154    fn rsqrt(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
155        let out = session.rsqrt_read(T::tensor_read(self))?;
156        into_typed_result("rsqrt", out)
157    }
158
159    fn compare(
160        &self,
161        rhs: &TypedTensor<T>,
162        dir: CompareDir,
163        session: &mut dyn Engine,
164    ) -> Result<TypedTensor<bool>> {
165        let (lhs, rhs) = broadcast_binary_in_read(self, rhs, session)?;
166        let out = session.compare_read(lhs.tensor_read(), rhs.tensor_read(), &dir)?;
167        into_typed_result("compare", out)
168    }
169
170    fn clamp(
171        &self,
172        lower: &TypedTensor<T>,
173        upper: &TypedTensor<T>,
174        session: &mut dyn Engine,
175    ) -> Result<TypedTensor<T>> {
176        let (input, lower, upper) = broadcast_ternary_in_read(self, lower, upper, session)?;
177        let out = session.clamp_read(
178            input.tensor_read(),
179            lower.tensor_read(),
180            upper.tensor_read(),
181        )?;
182        into_typed_result("clamp", out)
183    }
184
185    fn matmul(&self, rhs: &TypedTensor<T>, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
186        let config = matmul_config_for_shapes("matmul", self.shape(), rhs.shape())?;
187        let out = session.dot_general_read(T::tensor_read(self), T::tensor_read(rhs), &config)?;
188        into_typed_result("matmul", out)
189    }
190
191    fn reshape(&self, shape: &[usize], session: &mut dyn Engine) -> Result<TypedTensor<T>> {
192        let out = session.reshape_read(T::tensor_read(self), shape)?;
193        into_typed_result("reshape", out)
194    }
195
196    fn transpose(&self, perm: &[usize], session: &mut dyn Engine) -> Result<TypedTensor<T>> {
197        let out = session.transpose_read(T::tensor_read(self), perm)?;
198        into_typed_result("transpose", out)
199    }
200
201    fn broadcast_in_dim(
202        &self,
203        shape: &[usize],
204        dims: &[usize],
205        session: &mut dyn Engine,
206    ) -> Result<TypedTensor<T>> {
207        let out = session.broadcast_in_dim_read(T::tensor_read(self), shape, dims)?;
208        into_typed_result("broadcast_in_dim", out)
209    }
210
211    fn reduce_max(
212        &self,
213        axes: Option<&[usize]>,
214        session: &mut dyn Engine,
215    ) -> Result<TypedTensor<T>> {
216        let axes = crate::tensor::all_axes_if_none(self.shape().len(), axes);
217        let out = session.reduce_max_read(T::tensor_read(self), &axes)?;
218        into_typed_result("reduce_max", out)
219    }
220
221    fn reduce_min(
222        &self,
223        axes: Option<&[usize]>,
224        session: &mut dyn Engine,
225    ) -> Result<TypedTensor<T>> {
226        let axes = crate::tensor::all_axes_if_none(self.shape().len(), axes);
227        let out = session.reduce_min_read(T::tensor_read(self), &axes)?;
228        into_typed_result("reduce_min", out)
229    }
230
231    fn reduce_prod(
232        &self,
233        axes: Option<&[usize]>,
234        session: &mut dyn Engine,
235    ) -> Result<TypedTensor<T>> {
236        let axes = crate::tensor::all_axes_if_none(self.shape().len(), axes);
237        let out = session.reduce_prod_read(T::tensor_read(self), &axes)?;
238        into_typed_result("reduce_prod", out)
239    }
240
241    fn reduce_sum_squares(
242        &self,
243        axes: Option<&[usize]>,
244        session: &mut dyn Engine,
245    ) -> Result<TypedTensor<T>> {
246        let axes = crate::tensor::all_axes_if_none(self.shape().len(), axes);
247        let out = session.reduce_sum_squares_read(T::tensor_read(self), &axes)?;
248        into_typed_result("reduce_sum_squares", out)
249    }
250
251    fn dot_general(
252        &self,
253        rhs: &TypedTensor<T>,
254        config: DotGeneralConfig,
255        session: &mut dyn Engine,
256    ) -> Result<TypedTensor<T>> {
257        let out = session.dot_general_read(T::tensor_read(self), T::tensor_read(rhs), &config)?;
258        into_typed_result("dot_general", out)
259    }
260
261    fn dot_general_with_conj(
262        &self,
263        rhs: &TypedTensor<T>,
264        config: DotGeneralConfig,
265        lhs_conj: bool,
266        rhs_conj: bool,
267        session: &mut dyn Engine,
268    ) -> Result<TypedTensor<T>> {
269        let out = session.dot_general_with_conj_read(
270            T::tensor_read(self),
271            T::tensor_read(rhs),
272            &config,
273            lhs_conj,
274            rhs_conj,
275        )?;
276        into_typed_result("dot_general_with_conj", out)
277    }
278
279    fn scale_real(&self, factor: f64, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
280        let scalar = crate::scale::real_scale_scalar(T::dtype(), factor)?;
281        let scalar = session.upload_host_tensor(TensorRead::from_tensor(&scalar))?;
282        let scalar = into_typed_result::<T>("scale_real", scalar)?;
283        TypedTensorSessionOpsExt::mul(self, &scalar, session)
284    }
285
286    fn scale_complex(&self, factor: Complex64, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
287        let scalar = crate::scale::complex_scale_scalar(T::dtype(), factor)?;
288        let scalar = session.upload_host_tensor(TensorRead::from_tensor(&scalar))?;
289        let scalar = into_typed_result::<T>("scale_complex", scalar)?;
290        TypedTensorSessionOpsExt::mul(self, &scalar, session)
291    }
292
293    fn sigmoid(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
294        if let Some(out) =
295            session.fused_activation_read(ActivationOp::Sigmoid, T::tensor_read(self))?
296        {
297            return into_typed_result("sigmoid", out);
298        }
299        let out = run_session_composite(session, |ops| {
300            composite::sigmoid(ops, &typed_borrowed(self))
301        })?;
302        into_typed_result("sigmoid", out)
303    }
304
305    fn silu(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
306        if let Some(out) =
307            session.fused_activation_read(ActivationOp::Silu, T::tensor_read(self))?
308        {
309            return into_typed_result("silu", out);
310        }
311        let out =
312            run_session_composite(session, |ops| composite::silu(ops, &typed_borrowed(self)))?;
313        into_typed_result("silu", out)
314    }
315
316    fn softplus(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
317        if let Some(out) =
318            session.fused_activation_read(ActivationOp::Softplus, T::tensor_read(self))?
319        {
320            return into_typed_result("softplus", out);
321        }
322        let out = run_session_composite(session, |ops| {
323            composite::softplus(ops, &typed_borrowed(self))
324        })?;
325        into_typed_result("softplus", out)
326    }
327
328    fn gelu(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
329        if let Some(out) =
330            session.fused_activation_read(ActivationOp::Gelu, T::tensor_read(self))?
331        {
332            return into_typed_result("gelu", out);
333        }
334        let out =
335            run_session_composite(session, |ops| composite::gelu(ops, &typed_borrowed(self)))?;
336        into_typed_result("gelu", out)
337    }
338
339    fn gelu_tanh(&self, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
340        if let Some(out) =
341            session.fused_activation_read(ActivationOp::GeluTanh, T::tensor_read(self))?
342        {
343            return into_typed_result("gelu_tanh", out);
344        }
345        let out = run_session_composite(session, |ops| {
346            composite::gelu_tanh(ops, &typed_borrowed(self))
347        })?;
348        into_typed_result("gelu_tanh", out)
349    }
350
351    fn reduce_mean(
352        &self,
353        axes: Option<&[usize]>,
354        session: &mut dyn Engine,
355    ) -> Result<TypedTensor<T>> {
356        let out = run_session_composite(session, |ops| {
357            composite::reduce_mean(ops, &typed_borrowed(self), axes)
358        })?;
359        into_typed_result("reduce_mean", out)
360    }
361
362    fn softmax(&self, axis: usize, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
363        if let Some(output) = session.softmax_read(T::tensor_read(self), axis)? {
364            return into_typed_result("softmax", output);
365        }
366        let out = run_session_composite(session, |ops| {
367            composite::softmax(ops, &typed_borrowed(self), axis)
368        })?;
369        into_typed_result("softmax", out)
370    }
371
372    fn log_softmax(&self, axis: usize, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
373        let out = run_session_composite(session, |ops| {
374            composite::log_softmax(ops, &typed_borrowed(self), axis)
375        })?;
376        into_typed_result("log_softmax", out)
377    }
378
379    fn masked_softmax(
380        &self,
381        mask: &TypedTensor<bool>,
382        axis: usize,
383        session: &mut dyn Engine,
384    ) -> Result<TypedTensor<T>> {
385        let mask = typed_borrowed(mask);
386        let out = run_session_composite(session, |ops| {
387            composite::masked_softmax(ops, &typed_borrowed(self), &mask, axis)
388        })?;
389        into_typed_result("masked_softmax", out)
390    }
391
392    fn masked_log_softmax(
393        &self,
394        mask: &TypedTensor<bool>,
395        axis: usize,
396        session: &mut dyn Engine,
397    ) -> Result<TypedTensor<T>> {
398        let mask = typed_borrowed(mask);
399        let out = run_session_composite(session, |ops| {
400            composite::masked_log_softmax(ops, &typed_borrowed(self), &mask, axis)
401        })?;
402        into_typed_result("masked_log_softmax", out)
403    }
404
405    fn layer_norm(
406        &self,
407        axis: usize,
408        weight: Option<&TypedTensor<T>>,
409        bias: Option<&TypedTensor<T>>,
410        eps: f64,
411        session: &mut dyn Engine,
412    ) -> Result<TypedTensor<T>> {
413        let weight = weight.map(typed_borrowed);
414        let bias = bias.map(typed_borrowed);
415        let out = run_session_composite(session, |ops| {
416            composite::layer_norm(
417                ops,
418                &typed_borrowed(self),
419                axis,
420                weight.as_ref(),
421                bias.as_ref(),
422                eps,
423            )
424        })?;
425        into_typed_result("layer_norm", out)
426    }
427
428    fn rms_norm(
429        &self,
430        axis: usize,
431        weight: Option<&TypedTensor<T>>,
432        bias: Option<&TypedTensor<T>>,
433        eps: f64,
434        session: &mut dyn Engine,
435    ) -> Result<TypedTensor<T>> {
436        let weight = weight.map(typed_borrowed);
437        let bias = bias.map(typed_borrowed);
438        let out = run_session_composite(session, |ops| {
439            composite::rms_norm(
440                ops,
441                &typed_borrowed(self),
442                axis,
443                weight.as_ref(),
444                bias.as_ref(),
445                eps,
446            )
447        })?;
448        into_typed_result("rms_norm", out)
449    }
450
451    fn gather(
452        &self,
453        indices: &Tensor,
454        config: crate::GatherConfig,
455        session: &mut dyn Engine,
456    ) -> Result<TypedTensor<T>> {
457        let out = session.gather_read(
458            T::tensor_read(self),
459            TensorRead::from_tensor(indices),
460            &config,
461        )?;
462        into_typed_result("gather", out)
463    }
464
465    fn scatter(
466        &self,
467        indices: &Tensor,
468        updates: &Tensor,
469        config: crate::ScatterConfig,
470        session: &mut dyn Engine,
471    ) -> Result<TypedTensor<T>> {
472        let out = session.scatter_read(
473            T::tensor_read(self),
474            TensorRead::from_tensor(indices),
475            TensorRead::from_tensor(updates),
476            &config,
477        )?;
478        into_typed_result("scatter", out)
479    }
480
481    fn slice(
482        &self,
483        config: crate::SliceConfig,
484        session: &mut dyn Engine,
485    ) -> Result<TypedTensor<T>> {
486        let out = session.slice_read(T::tensor_read(self), &config)?;
487        into_typed_result("slice", out)
488    }
489
490    fn dynamic_slice(
491        &self,
492        starts: &Tensor,
493        sizes: &[usize],
494        session: &mut dyn Engine,
495    ) -> Result<TypedTensor<T>> {
496        let out = session.dynamic_slice_read(
497            T::tensor_read(self),
498            TensorRead::from_tensor(starts),
499            sizes,
500        )?;
501        into_typed_result("dynamic_slice", out)
502    }
503
504    fn dynamic_update_slice(
505        &self,
506        update: &TypedTensor<T>,
507        starts: &Tensor,
508        session: &mut dyn Engine,
509    ) -> Result<TypedTensor<T>> {
510        let out = session.dynamic_update_slice_read(
511            T::tensor_read(self),
512            T::tensor_read(update),
513            TensorRead::from_tensor(starts),
514        )?;
515        into_typed_result("dynamic_update_slice", out)
516    }
517
518    fn pad(&self, config: crate::PadConfig, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
519        let out = session.pad_read(T::tensor_read(self), &config)?;
520        into_typed_result("pad", out)
521    }
522
523    fn concatenate(
524        inputs: &[&TypedTensor<T>],
525        axis: usize,
526        session: &mut dyn Engine,
527    ) -> Result<TypedTensor<T>> {
528        let reads: Vec<TensorRead<'_>> = inputs.iter().map(|input| T::tensor_read(input)).collect();
529        let out = session.concatenate_read(&reads, axis)?;
530        into_typed_result("concatenate", out)
531    }
532
533    fn reverse(&self, axes: &[usize], session: &mut dyn Engine) -> Result<TypedTensor<T>> {
534        let out = session.reverse_read(T::tensor_read(self), axes)?;
535        into_typed_result("reverse", out)
536    }
537
538    fn tril(&self, k: i64, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
539        let out = session.tril_read(T::tensor_read(self), k)?;
540        into_typed_result("tril", out)
541    }
542
543    fn triu(&self, k: i64, session: &mut dyn Engine) -> Result<TypedTensor<T>> {
544        let out = session.triu_read(T::tensor_read(self), k)?;
545        into_typed_result("triu", out)
546    }
547
548    fn extract_diag(
549        &self,
550        axis_a: usize,
551        axis_b: usize,
552        session: &mut dyn Engine,
553    ) -> Result<TypedTensor<T>> {
554        let out = session.extract_diagonal_read(T::tensor_read(self), axis_a, axis_b)?;
555        into_typed_result("extract_diag", out)
556    }
557
558    fn embed_diag(
559        &self,
560        axis_a: usize,
561        axis_b: usize,
562        session: &mut dyn Engine,
563    ) -> Result<TypedTensor<T>> {
564        let out = session.embed_diagonal_read(T::tensor_read(self), axis_a, axis_b)?;
565        into_typed_result("embed_diag", out)
566    }
567
568    fn take_along_axis(
569        &self,
570        indices: &Tensor,
571        axis: usize,
572        session: &mut dyn Engine,
573    ) -> Result<TypedTensor<T>> {
574        let indices = borrowed(indices);
575        let out = run_session_composite(session, |ops| {
576            composite::take_along_axis(ops, &typed_borrowed(self), &indices, axis)
577        })?;
578        into_typed_result("take_along_axis", out)
579    }
580}
581
582impl TypedTensorMaskSessionOpsExt for TypedTensor<bool> {
583    fn where_select<U: TensorScalar>(
584        &self,
585        on_true: &TypedTensor<U>,
586        on_false: &TypedTensor<U>,
587        session: &mut dyn Engine,
588    ) -> Result<TypedTensor<U>> {
589        let (condition, on_true, on_false) =
590            broadcast_ternary_in_read(self, on_true, on_false, session)?;
591        let out = session.select_read(
592            condition.tensor_read(),
593            on_true.tensor_read(),
594            on_false.tensor_read(),
595        )?;
596        into_typed_result("where_select", out)
597    }
598}
599
600// INVARIANT: this private adapter keeps borrowed reads borrowed, owns the
601// explicit fallback tensor, and owns an intermediate whose layout is a
602// metadata-only view of that tensor; it is never exposed or cloned.
603#[allow(clippy::large_enum_variant)]
604pub(crate) enum ReadInput<'a> {
605    Borrowed(TensorRead<'a>),
606    Owned(Tensor),
607    Value(TensorValue),
608}
609
610impl ReadInput<'_> {
611    pub(crate) fn tensor_read(&self) -> TensorRead<'_> {
612        match self {
613            Self::Borrowed(read) => read.clone(),
614            Self::Owned(tensor) => TensorRead::from_tensor(tensor),
615            Self::Value(value) => value.tensor_read(),
616        }
617    }
618}
619
620pub(crate) fn broadcast_to_in_read<'a>(
621    input: TensorRead<'a>,
622    target_shape: &[usize],
623    session: &mut dyn Engine,
624) -> Result<ReadInput<'a>> {
625    if input.shape() == target_shape {
626        return Ok(ReadInput::Borrowed(input));
627    }
628
629    // The eager surface aligns the operand's axes with the trailing output axes
630    // and reads expanding axes through zero strides, so an implicit broadcast is
631    // metadata. The traced/AD lowering instead removes expanding singleton axes
632    // before broadcasting; that plan needs a materializing reshape and is not
633    // used here (#2143 item 1).
634    let dims = broadcast_aligned_dims(input.shape(), target_shape).map_err(broadcast_error)?;
635
636    // A caller-owned payload has no borrowed view, so it keeps the session's
637    // typed refusal path rather than a view.
638    if let DType::External(_) = input.dtype() {
639        let out = session.broadcast_in_dim_read(input, target_shape, &dims)?;
640        return Ok(ReadInput::Owned(out));
641    }
642
643    Ok(ReadInput::Borrowed(
644        input.broadcast_in_dim_view(target_shape, &dims)?,
645    ))
646}
647
648/// The trailing-axis mapping of `input` onto `output`, or the same broadcast
649/// failure the traced/AD plan reports.
650fn broadcast_aligned_dims(
651    input: &[usize],
652    output: &[usize],
653) -> std::result::Result<TinyVec<[usize; 8]>, BroadcastError> {
654    if input.len() > output.len() {
655        return Err(BroadcastError::RankTooLarge {
656            input: input.to_vec(),
657            output: output.to_vec(),
658        });
659    }
660    let rank_diff = output.len() - input.len();
661    for (src_axis, &src_dim) in input.iter().enumerate() {
662        let dst_dim = output[src_axis + rank_diff];
663        if src_dim != dst_dim && src_dim != 1 {
664            return Err(BroadcastError::IncompatibleInput {
665                input: input.to_vec(),
666                output: output.to_vec(),
667            });
668        }
669    }
670    Ok((0..input.len()).map(|axis| axis + rank_diff).collect())
671}
672
673fn broadcast_binary_in_read<'a, T: TensorScalar>(
674    lhs: &'a TypedTensor<T>,
675    rhs: &'a TypedTensor<T>,
676    session: &mut dyn Engine,
677) -> Result<(ReadInput<'a>, ReadInput<'a>)> {
678    let shape = broadcast_shape(lhs.shape(), rhs.shape()).map_err(broadcast_error)?;
679    Ok((
680        broadcast_to_in_read(T::tensor_read(lhs), &shape, session)?,
681        broadcast_to_in_read(T::tensor_read(rhs), &shape, session)?,
682    ))
683}
684
685fn broadcast_ternary_in_read<'a, C: TensorScalar, T: TensorScalar>(
686    first: &'a TypedTensor<C>,
687    second: &'a TypedTensor<T>,
688    third: &'a TypedTensor<T>,
689    session: &mut dyn Engine,
690) -> Result<(ReadInput<'a>, ReadInput<'a>, ReadInput<'a>)> {
691    let shape = broadcast_shapes([first.shape(), second.shape(), third.shape()])
692        .map_err(broadcast_error)?;
693    Ok((
694        broadcast_to_in_read(C::tensor_read(first), &shape, session)?,
695        broadcast_to_in_read(T::tensor_read(second), &shape, session)?,
696        broadcast_to_in_read(T::tensor_read(third), &shape, session)?,
697    ))
698}
699
700pub(crate) fn broadcast_error(err: BroadcastError) -> Error {
701    match err {
702        BroadcastError::IncompatibleBinary { lhs, rhs } => {
703            Error::shape_mismatch("broadcast", lhs, rhs)
704        }
705        BroadcastError::IncompatibleInput { input, output } => {
706            Error::shape_mismatch("broadcast", input, output)
707        }
708        BroadcastError::RankTooLarge { input, output } => {
709            Error::rank_mismatch("broadcast", output.len(), input.len())
710        }
711    }
712}
713
714fn into_typed_result<T: TensorScalar>(op: &'static str, tensor: Tensor) -> Result<TypedTensor<T>> {
715    let actual = tensor.dtype();
716    T::into_typed(tensor).map_err(|_| {
717        Error::validation(
718            op,
719            ValidationError::DTypeMismatch {
720                expected: T::dtype(),
721                actual,
722            },
723        )
724    })
725}
tenferro_runtime/typed_tensor.rs

tenferro_runtime/
typed_tensor.rs
