1use 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
206fn 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 #[doc = concat!("let y = x.", stringify!($name), "(&mut backend).unwrap();")]
241 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 #[doc = concat!("let z = x.", stringify!($name), "(&y, &mut backend).unwrap();")]
265 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
324fn 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
347fn 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
375fn 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
408fn 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
438fn 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
464fn 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
490fn 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
512fn 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
534fn 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}