1use std::ops::{Add, Div, Mul, Neg, Rem as StdRem, Sub};
2use std::sync::Arc;
3
4use num_complex::Complex;
5use num_traits::{One, Zero};
6use strided_kernel::{
7 batched_outer_product_into, broadcast_mul_into, fused_elementwise_into, map_into, mul_into,
8 reduce, zip_map2_into, zip_map3_into, FusedInst, FusedOp, FusedPlan, FusedScalar, StridedView,
9};
10
11use crate::buffer_pool::{BufferPool, PoolScalar};
12use crate::ConjElem;
13use tenferro_tensor::backend::{
14 ElementwiseFusionInputView, ElementwiseFusionOp, ElementwiseFusionPlan,
15};
16use tenferro_tensor::{
17 col_major_strides, CompareDir, DType, Tensor, TensorOwnedView, TensorRank, TensorRead,
18 TensorScalar, TensorValue, TensorView, TypedTensor, TypedTensorView,
19};
20
21use super::{
22 tensor_from_array, typed_array_uninit_from_pool, typed_host_data, typed_view,
23 typed_view_from_view,
24};
25
26macro_rules! dispatch_ternary_result_with_pool {
27 ($op:literal, $a:expr, $b:expr, $c:expr, |$x:ident, $y:ident, $z:ident| $body:expr) => {
28 match ($a, $b, $c) {
29 (Tensor::F32($x), Tensor::F32($y), Tensor::F32($z)) => Ok(Tensor::F32($body?)),
30 (Tensor::F64($x), Tensor::F64($y), Tensor::F64($z)) => Ok(Tensor::F64($body?)),
31 _ => Err(crate::Error::backend_failure($op, "dtype mismatch")),
32 }
33 };
34}
35
36fn dtype_pair_error(op: &'static str, lhs: DType, rhs: DType) -> crate::Error {
37 if lhs == rhs {
38 crate::Error::backend_failure(op, format!("unsupported dtype {lhs:?}"))
39 } else {
40 crate::Error::DTypeMismatch { op, lhs, rhs }
41 }
42}
43
44fn tensor_pair_error(op: &'static str, lhs: &Tensor, rhs: &Tensor) -> crate::Error {
45 dtype_pair_error(op, lhs.dtype(), rhs.dtype())
46}
47
48fn read_pair_error(op: &'static str, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Error {
49 dtype_pair_error(op, lhs.dtype(), rhs.dtype())
50}
51
52fn is_complex_dtype(dtype: DType) -> bool {
53 matches!(dtype, DType::C32 | DType::C64)
54}
55
56fn ordered_complex_error(op: &'static str) -> crate::Error {
57 crate::Error::InvalidConfig {
58 op,
59 message: "complex tensors do not have a total order; compute abs/norm explicitly before ordered operations".into(),
60 }
61}
62
63fn reject_complex_ordered_dtypes(op: &'static str, dtypes: &[DType]) -> crate::Result<()> {
64 if dtypes.iter().copied().any(is_complex_dtype) {
65 return Err(ordered_complex_error(op));
66 }
67 Ok(())
68}
69
70const ELEMENTWISE_FUSION_OP: &str = "execute_elementwise_fusion";
71const ELEMENTWISE_FUSION_MIN_ELEMENTS: usize = 16 * 1024;
72
73fn validate_elementwise_fusion_inputs(
74 inputs: &[&Tensor],
75 plan: &ElementwiseFusionPlan,
76) -> crate::Result<bool> {
77 if inputs.len() != plan.input_count() {
78 return Err(crate::Error::backend_failure(
79 ELEMENTWISE_FUSION_OP,
80 format!(
81 "plan expects {} inputs but backend received {}",
82 plan.input_count(),
83 inputs.len()
84 ),
85 ));
86 }
87 if plan.input_views().len() != plan.input_count() {
88 return Err(crate::Error::backend_failure(
89 ELEMENTWISE_FUSION_OP,
90 format!(
91 "plan has {} input views for {} inputs",
92 plan.input_views().len(),
93 plan.input_count()
94 ),
95 ));
96 }
97 if plan.outputs().is_empty() {
98 return Ok(false);
99 }
100 for input in inputs {
101 if input.dtype() != plan.dtype() {
102 return Err(crate::Error::DTypeMismatch {
103 op: ELEMENTWISE_FUSION_OP,
104 lhs: input.dtype(),
105 rhs: plan.dtype(),
106 });
107 }
108 }
109 Ok(true)
110}
111
112fn strided_fused_op(op: ElementwiseFusionOp) -> FusedOp {
113 match op {
114 ElementwiseFusionOp::Add => FusedOp::Add,
115 ElementwiseFusionOp::Multiply => FusedOp::Multiply,
116 ElementwiseFusionOp::Negate => FusedOp::Negate,
117 ElementwiseFusionOp::Conj => FusedOp::Conj,
118 ElementwiseFusionOp::Divide => FusedOp::Divide,
119 ElementwiseFusionOp::Abs => FusedOp::Abs,
120 ElementwiseFusionOp::Maximum => FusedOp::Maximum,
121 ElementwiseFusionOp::Minimum => FusedOp::Minimum,
122 ElementwiseFusionOp::Clamp => FusedOp::Clamp,
123 ElementwiseFusionOp::Exp => FusedOp::Exp,
124 ElementwiseFusionOp::Log => FusedOp::Log,
125 ElementwiseFusionOp::Sin => FusedOp::Sin,
126 ElementwiseFusionOp::Cos => FusedOp::Cos,
127 ElementwiseFusionOp::Tanh => FusedOp::Tanh,
128 ElementwiseFusionOp::Sqrt => FusedOp::Sqrt,
129 ElementwiseFusionOp::Rsqrt => FusedOp::Rsqrt,
130 ElementwiseFusionOp::Pow => FusedOp::Pow,
131 ElementwiseFusionOp::Expm1 => FusedOp::Expm1,
132 ElementwiseFusionOp::Log1p => FusedOp::Log1p,
133 ElementwiseFusionOp::Remainder => {
134 unreachable!("remainder must be filtered before CPU elementwise fusion")
135 }
136 }
137}
138
139fn plan_uses_unfused_op(plan: &ElementwiseFusionPlan) -> bool {
140 plan.ops()
141 .iter()
142 .any(|inst| inst.op() == ElementwiseFusionOp::Remainder)
143}
144
145fn plan_uses_ordered_op(plan: &ElementwiseFusionPlan) -> bool {
146 plan.ops().iter().any(|inst| {
147 matches!(
148 inst.op(),
149 ElementwiseFusionOp::Maximum
150 | ElementwiseFusionOp::Minimum
151 | ElementwiseFusionOp::Clamp
152 )
153 })
154}
155
156fn should_defer_to_broadcast_multiply_special_case(plan: &ElementwiseFusionPlan) -> bool {
157 !plan.input_views().iter().all(|view| view.is_identity())
158 && plan.ops().len() == 1
159 && plan.outputs() == [plan.input_count()]
160 && plan.ops()[0].op() == ElementwiseFusionOp::Multiply
161}
162
163fn strided_fused_plan(plan: &ElementwiseFusionPlan) -> FusedPlan {
164 FusedPlan {
165 input_count: plan.input_count(),
166 outputs: plan.outputs().to_vec(),
167 ops: plan
168 .ops()
169 .iter()
170 .map(|inst| FusedInst {
171 op: strided_fused_op(inst.op()),
172 inputs: inst.inputs().to_vec(),
173 })
174 .collect(),
175 }
176}
177
178pub(crate) trait Tier2Elem: Copy + Clone + One + Zero + Send + Sync {
179 fn abs_elem(self) -> Self;
180 fn sign_elem(self) -> Self;
181}
182
183pub(crate) trait OrderedElem: Copy + Clone + Send + Sync {
186 fn max_elem(self, other: Self) -> Self;
187 fn min_elem(self, other: Self) -> Self;
188}
189
190pub(crate) trait CompareElem: Copy + Send + Sync {
191 fn compare_elem(self, other: Self, dir: &CompareDir) -> bool;
192}
193
194trait WrappingIntegerElem:
195 Copy + PoolScalar + TensorScalar + Zero + PartialEq + Eq + Send + Sync + 'static
196{
197 fn wrapping_add_elem(self, other: Self) -> Self;
198 fn wrapping_sub_elem(self, other: Self) -> Self;
199 fn wrapping_mul_elem(self, other: Self) -> Self;
200 fn wrapping_div_elem(self, other: Self) -> Self;
201 fn wrapping_rem_elem(self, other: Self) -> Self;
202 fn wrapping_neg_elem(self) -> Self;
203 fn wrapping_abs_elem(self) -> Self;
204 fn signum_elem(self) -> Self;
205}
206
207macro_rules! impl_tier2_elem_real {
208 ($ty:ty) => {
209 impl Tier2Elem for $ty {
210 fn abs_elem(self) -> Self {
211 self.abs()
212 }
213
214 fn sign_elem(self) -> Self {
215 if self == Self::zero() {
216 Self::zero()
217 } else {
218 self.signum()
219 }
220 }
221 }
222
223 impl OrderedElem for $ty {
224 fn max_elem(self, other: Self) -> Self {
225 if self.is_nan() || other.is_nan() {
226 <$ty>::NAN
227 } else if self >= other {
228 self
229 } else {
230 other
231 }
232 }
233
234 fn min_elem(self, other: Self) -> Self {
235 if self.is_nan() || other.is_nan() {
236 <$ty>::NAN
237 } else if self <= other {
238 self
239 } else {
240 other
241 }
242 }
243 }
244
245 impl CompareElem for $ty {
246 fn compare_elem(self, other: Self, dir: &CompareDir) -> bool {
247 match dir {
248 CompareDir::Eq => self == other,
249 CompareDir::Lt => self < other,
250 CompareDir::Le => self <= other,
251 CompareDir::Gt => self > other,
252 CompareDir::Ge => self >= other,
253 }
254 }
255 }
256 };
257}
258
259macro_rules! impl_tier2_elem_complex {
260 ($real:ty) => {
261 impl Tier2Elem for Complex<$real> {
262 fn abs_elem(self) -> Self {
263 Self::new(self.norm(), <$real>::zero())
264 }
265
266 fn sign_elem(self) -> Self {
267 if self.is_zero() {
268 Self::zero()
269 } else {
270 self / self.abs_elem()
271 }
272 }
273 }
274 };
275}
276
277impl_tier2_elem_real!(f32);
278impl_tier2_elem_real!(f64);
279impl_tier2_elem_complex!(f32);
280impl_tier2_elem_complex!(f64);
281
282macro_rules! impl_compare_elem_ord {
283 ($ty:ty) => {
284 impl CompareElem for $ty {
285 fn compare_elem(self, other: Self, dir: &CompareDir) -> bool {
286 match dir {
287 CompareDir::Eq => self == other,
288 CompareDir::Lt => self < other,
289 CompareDir::Le => self <= other,
290 CompareDir::Gt => self > other,
291 CompareDir::Ge => self >= other,
292 }
293 }
294 }
295 };
296}
297
298impl_compare_elem_ord!(i32);
299impl_compare_elem_ord!(i64);
300impl_compare_elem_ord!(bool);
301
302macro_rules! impl_ordered_elem_ord {
303 ($ty:ty) => {
304 impl OrderedElem for $ty {
305 fn max_elem(self, other: Self) -> Self {
306 self.max(other)
307 }
308
309 fn min_elem(self, other: Self) -> Self {
310 self.min(other)
311 }
312 }
313 };
314}
315
316impl_ordered_elem_ord!(i32);
317impl_ordered_elem_ord!(i64);
318
319macro_rules! impl_wrapping_integer_elem {
320 ($ty:ty) => {
321 impl WrappingIntegerElem for $ty {
322 fn wrapping_add_elem(self, other: Self) -> Self {
323 self.wrapping_add(other)
324 }
325
326 fn wrapping_sub_elem(self, other: Self) -> Self {
327 self.wrapping_sub(other)
328 }
329
330 fn wrapping_mul_elem(self, other: Self) -> Self {
331 self.wrapping_mul(other)
332 }
333
334 fn wrapping_div_elem(self, other: Self) -> Self {
335 self.wrapping_div(other)
336 }
337
338 fn wrapping_rem_elem(self, other: Self) -> Self {
339 self.wrapping_rem(other)
340 }
341
342 fn wrapping_neg_elem(self) -> Self {
343 self.wrapping_neg()
344 }
345
346 fn wrapping_abs_elem(self) -> Self {
347 self.wrapping_abs()
348 }
349
350 fn signum_elem(self) -> Self {
351 self.signum()
352 }
353 }
354 };
355}
356
357impl_wrapping_integer_elem!(i32);
358impl_wrapping_integer_elem!(i64);
359
360fn strided_view_contains<T>(
361 op: &'static str,
362 view: &StridedView<'_, T>,
363 pred: impl Fn(T) -> bool + Copy + Sync,
364) -> crate::Result<bool>
365where
366 T: Copy + Send + Sync,
367{
368 reduce(view, pred, |lhs, rhs| lhs || rhs, false)
369 .map_err(|err| crate::Error::backend_failure(op, err))
370}
371
372fn ensure_no_zero_divisor<T>(op: &'static str, rhs: &StridedView<'_, T>) -> crate::Result<()>
373where
374 T: WrappingIntegerElem,
375{
376 if strided_view_contains(op, rhs, |value| value == T::zero())? {
377 return Err(crate::Error::division_by_zero(op, T::dtype()));
378 }
379 Ok(())
380}
381
382fn complex_scalar_tensor<T>(scalar: T) -> crate::Result<TypedTensor<Complex<T>>>
383where
384 T: Copy + Clone + Zero,
385{
386 TypedTensor::from_vec_col_major(vec![], vec![Complex::new(scalar, T::zero())])
387}
388
389fn complex_scalar_tensor_from_tensor<T>(
390 input: &TypedTensor<T>,
391) -> crate::Result<TypedTensor<Complex<T>>>
392where
393 T: Copy + Clone + Zero,
394{
395 complex_scalar_tensor(typed_host_data("add", input)?[0])
396}
397
398fn complex_scalar_tensor_from_view<T, R>(
399 input: &TypedTensorView<'_, T, R>,
400) -> crate::Result<TypedTensor<Complex<T>>>
401where
402 T: Copy + Clone + Zero + 'static,
403 R: TensorRank,
404{
405 complex_scalar_tensor(typed_view_from_view("add", input)?.get(&[]))
406}
407
408fn with_local_pool<T>(f: impl FnOnce(&mut BufferPool) -> T) -> T {
409 let mut buffers = BufferPool::new();
410 f(&mut buffers)
411}
412
413pub fn add(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
428 with_local_pool(|buffers| add_with_pool(buffers, lhs, rhs))
429}
430
431pub(crate) fn add_with_pool(
432 buffers: &mut BufferPool,
433 lhs: &Tensor,
434 rhs: &Tensor,
435) -> crate::Result<Tensor> {
436 match (lhs, rhs) {
437 (Tensor::F32(a), Tensor::F32(b)) => Ok(Tensor::F32(typed_add_with_pool(buffers, a, b)?)),
438 (Tensor::F64(a), Tensor::F64(b)) => Ok(Tensor::F64(typed_add_with_pool(buffers, a, b)?)),
439 (Tensor::I32(a), Tensor::I32(b)) => {
440 Ok(Tensor::I32(typed_wrapping_add_with_pool(buffers, a, b)?))
441 }
442 (Tensor::I64(a), Tensor::I64(b)) => {
443 Ok(Tensor::I64(typed_wrapping_add_with_pool(buffers, a, b)?))
444 }
445 (Tensor::C32(a), Tensor::C32(b)) => Ok(Tensor::C32(typed_add_with_pool(buffers, a, b)?)),
446 (Tensor::C64(a), Tensor::C64(b)) => Ok(Tensor::C64(typed_add_with_pool(buffers, a, b)?)),
447 (Tensor::F32(a), Tensor::C32(b)) if a.shape().is_empty() => {
448 let scalar = complex_scalar_tensor(typed_host_data("add", a)?[0])?;
449 Ok(Tensor::C32(typed_add_with_pool(buffers, &scalar, b)?))
450 }
451 (Tensor::C32(a), Tensor::F32(b)) if b.shape().is_empty() => {
452 let scalar = complex_scalar_tensor(typed_host_data("add", b)?[0])?;
453 Ok(Tensor::C32(typed_add_with_pool(buffers, a, &scalar)?))
454 }
455 (Tensor::F64(a), Tensor::C64(b)) if a.shape().is_empty() => {
456 let scalar = complex_scalar_tensor(typed_host_data("add", a)?[0])?;
457 Ok(Tensor::C64(typed_add_with_pool(buffers, &scalar, b)?))
458 }
459 (Tensor::C64(a), Tensor::F64(b)) if b.shape().is_empty() => {
460 let scalar = complex_scalar_tensor(typed_host_data("add", b)?[0])?;
461 Ok(Tensor::C64(typed_add_with_pool(buffers, a, &scalar)?))
462 }
463 _ => Err(tensor_pair_error("add", lhs, rhs)),
464 }
465}
466
467pub(crate) fn add_read_with_pool(
468 buffers: &mut BufferPool,
469 lhs: TensorRead<'_>,
470 rhs: TensorRead<'_>,
471) -> crate::Result<Tensor> {
472 if let (TensorRead::Tensor(lhs), TensorRead::Tensor(rhs)) = (&lhs, &rhs) {
473 return add_with_pool(buffers, lhs, rhs);
474 }
475
476 macro_rules! dispatch {
477 ($variant:ident, $func:ident) => {
478 match (&lhs, &rhs) {
479 (
480 TensorRead::Tensor(Tensor::$variant(a)),
481 TensorRead::View(TensorView::$variant(b)),
482 ) => {
483 let a = a.as_view();
484 return Ok(Tensor::$variant($func(buffers, &a, b)?));
485 }
486 (
487 TensorRead::View(TensorView::$variant(a)),
488 TensorRead::Tensor(Tensor::$variant(b)),
489 ) => {
490 let b = b.as_view();
491 return Ok(Tensor::$variant($func(buffers, a, &b)?));
492 }
493 (
494 TensorRead::View(TensorView::$variant(a)),
495 TensorRead::View(TensorView::$variant(b)),
496 ) => {
497 return Ok(Tensor::$variant($func(buffers, a, b)?));
498 }
499 _ => {}
500 }
501 };
502 }
503
504 macro_rules! dispatch_real_complex_scalar {
505 ($real_variant:ident, $complex_variant:ident) => {
506 match (&lhs, &rhs) {
507 (
508 TensorRead::Tensor(Tensor::$real_variant(real)),
509 TensorRead::View(TensorView::$complex_variant(complex)),
510 ) if real.shape().is_empty() => {
511 let scalar = complex_scalar_tensor_from_tensor(real)?;
512 let scalar = scalar.as_view();
513 return Ok(Tensor::$complex_variant(typed_add_view_with_pool(
514 buffers, &scalar, complex,
515 )?));
516 }
517 (
518 TensorRead::View(TensorView::$real_variant(real)),
519 TensorRead::Tensor(Tensor::$complex_variant(complex)),
520 ) if real.shape().is_empty() => {
521 let scalar = complex_scalar_tensor_from_view(real)?;
522 let scalar = scalar.as_view();
523 let complex = complex.as_view();
524 return Ok(Tensor::$complex_variant(typed_add_view_with_pool(
525 buffers, &scalar, &complex,
526 )?));
527 }
528 (
529 TensorRead::View(TensorView::$real_variant(real)),
530 TensorRead::View(TensorView::$complex_variant(complex)),
531 ) if real.shape().is_empty() => {
532 let scalar = complex_scalar_tensor_from_view(real)?;
533 let scalar = scalar.as_view();
534 return Ok(Tensor::$complex_variant(typed_add_view_with_pool(
535 buffers, &scalar, complex,
536 )?));
537 }
538 (
539 TensorRead::Tensor(Tensor::$complex_variant(complex)),
540 TensorRead::View(TensorView::$real_variant(real)),
541 ) if real.shape().is_empty() => {
542 let complex = complex.as_view();
543 let scalar = complex_scalar_tensor_from_view(real)?;
544 let scalar = scalar.as_view();
545 return Ok(Tensor::$complex_variant(typed_add_view_with_pool(
546 buffers, &complex, &scalar,
547 )?));
548 }
549 (
550 TensorRead::View(TensorView::$complex_variant(complex)),
551 TensorRead::Tensor(Tensor::$real_variant(real)),
552 ) if real.shape().is_empty() => {
553 let scalar = complex_scalar_tensor_from_tensor(real)?;
554 let scalar = scalar.as_view();
555 return Ok(Tensor::$complex_variant(typed_add_view_with_pool(
556 buffers, complex, &scalar,
557 )?));
558 }
559 (
560 TensorRead::View(TensorView::$complex_variant(complex)),
561 TensorRead::View(TensorView::$real_variant(real)),
562 ) if real.shape().is_empty() => {
563 let scalar = complex_scalar_tensor_from_view(real)?;
564 let scalar = scalar.as_view();
565 return Ok(Tensor::$complex_variant(typed_add_view_with_pool(
566 buffers, complex, &scalar,
567 )?));
568 }
569 _ => {}
570 }
571 };
572 }
573
574 dispatch_real_complex_scalar!(F32, C32);
575 dispatch_real_complex_scalar!(F64, C64);
576
577 dispatch!(F32, typed_add_view_with_pool);
578 dispatch!(F64, typed_add_view_with_pool);
579 dispatch!(I32, typed_wrapping_add_view_with_pool);
580 dispatch!(I64, typed_wrapping_add_view_with_pool);
581 dispatch!(C32, typed_add_view_with_pool);
582 dispatch!(C64, typed_add_view_with_pool);
583
584 Err(read_pair_error("add", lhs, rhs))
585}
586
587pub fn sub(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
602 with_local_pool(|buffers| sub_with_pool(buffers, lhs, rhs))
603}
604
605pub(crate) fn sub_with_pool(
606 buffers: &mut BufferPool,
607 lhs: &Tensor,
608 rhs: &Tensor,
609) -> crate::Result<Tensor> {
610 match (lhs, rhs) {
611 (Tensor::F32(a), Tensor::F32(b)) => Ok(Tensor::F32(typed_sub_with_pool(buffers, a, b)?)),
612 (Tensor::F64(a), Tensor::F64(b)) => Ok(Tensor::F64(typed_sub_with_pool(buffers, a, b)?)),
613 (Tensor::I32(a), Tensor::I32(b)) => {
614 Ok(Tensor::I32(typed_wrapping_sub_with_pool(buffers, a, b)?))
615 }
616 (Tensor::I64(a), Tensor::I64(b)) => {
617 Ok(Tensor::I64(typed_wrapping_sub_with_pool(buffers, a, b)?))
618 }
619 (Tensor::C32(a), Tensor::C32(b)) => Ok(Tensor::C32(typed_sub_with_pool(buffers, a, b)?)),
620 (Tensor::C64(a), Tensor::C64(b)) => Ok(Tensor::C64(typed_sub_with_pool(buffers, a, b)?)),
621 (Tensor::F32(a), Tensor::C32(b)) if a.shape().is_empty() => {
622 let scalar = complex_scalar_tensor(typed_host_data("sub", a)?[0])?;
623 Ok(Tensor::C32(typed_sub_with_pool(buffers, &scalar, b)?))
624 }
625 (Tensor::C32(a), Tensor::F32(b)) if b.shape().is_empty() => {
626 let scalar = complex_scalar_tensor(typed_host_data("sub", b)?[0])?;
627 Ok(Tensor::C32(typed_sub_with_pool(buffers, a, &scalar)?))
628 }
629 (Tensor::F64(a), Tensor::C64(b)) if a.shape().is_empty() => {
630 let scalar = complex_scalar_tensor(typed_host_data("sub", a)?[0])?;
631 Ok(Tensor::C64(typed_sub_with_pool(buffers, &scalar, b)?))
632 }
633 (Tensor::C64(a), Tensor::F64(b)) if b.shape().is_empty() => {
634 let scalar = complex_scalar_tensor(typed_host_data("sub", b)?[0])?;
635 Ok(Tensor::C64(typed_sub_with_pool(buffers, a, &scalar)?))
636 }
637 _ => Err(tensor_pair_error("sub", lhs, rhs)),
638 }
639}
640
641pub(crate) fn sub_read_with_pool(
642 buffers: &mut BufferPool,
643 lhs: TensorRead<'_>,
644 rhs: TensorRead<'_>,
645) -> crate::Result<Tensor> {
646 if let (TensorRead::Tensor(lhs), TensorRead::Tensor(rhs)) = (&lhs, &rhs) {
647 return sub_with_pool(buffers, lhs, rhs);
648 }
649
650 macro_rules! dispatch {
651 ($variant:ident, $func:ident) => {
652 match (&lhs, &rhs) {
653 (
654 TensorRead::Tensor(Tensor::$variant(a)),
655 TensorRead::View(TensorView::$variant(b)),
656 ) => {
657 let a = a.as_view();
658 return Ok(Tensor::$variant($func(buffers, &a, b)?));
659 }
660 (
661 TensorRead::View(TensorView::$variant(a)),
662 TensorRead::Tensor(Tensor::$variant(b)),
663 ) => {
664 let b = b.as_view();
665 return Ok(Tensor::$variant($func(buffers, a, &b)?));
666 }
667 (
668 TensorRead::View(TensorView::$variant(a)),
669 TensorRead::View(TensorView::$variant(b)),
670 ) => {
671 return Ok(Tensor::$variant($func(buffers, a, b)?));
672 }
673 _ => {}
674 }
675 };
676 }
677
678 macro_rules! dispatch_real_complex_scalar {
679 ($real_variant:ident, $complex_variant:ident) => {
680 match (&lhs, &rhs) {
681 (
682 TensorRead::Tensor(Tensor::$real_variant(real)),
683 TensorRead::View(TensorView::$complex_variant(complex)),
684 ) if real.shape().is_empty() => {
685 let scalar = complex_scalar_tensor_from_tensor(real)?;
686 let scalar = scalar.as_view();
687 return Ok(Tensor::$complex_variant(typed_sub_view_with_pool(
688 buffers, &scalar, complex,
689 )?));
690 }
691 (
692 TensorRead::View(TensorView::$real_variant(real)),
693 TensorRead::Tensor(Tensor::$complex_variant(complex)),
694 ) if real.shape().is_empty() => {
695 let scalar = complex_scalar_tensor_from_view(real)?;
696 let scalar = scalar.as_view();
697 let complex = complex.as_view();
698 return Ok(Tensor::$complex_variant(typed_sub_view_with_pool(
699 buffers, &scalar, &complex,
700 )?));
701 }
702 (
703 TensorRead::View(TensorView::$real_variant(real)),
704 TensorRead::View(TensorView::$complex_variant(complex)),
705 ) if real.shape().is_empty() => {
706 let scalar = complex_scalar_tensor_from_view(real)?;
707 let scalar = scalar.as_view();
708 return Ok(Tensor::$complex_variant(typed_sub_view_with_pool(
709 buffers, &scalar, complex,
710 )?));
711 }
712 (
713 TensorRead::Tensor(Tensor::$complex_variant(complex)),
714 TensorRead::View(TensorView::$real_variant(real)),
715 ) if real.shape().is_empty() => {
716 let complex = complex.as_view();
717 let scalar = complex_scalar_tensor_from_view(real)?;
718 let scalar = scalar.as_view();
719 return Ok(Tensor::$complex_variant(typed_sub_view_with_pool(
720 buffers, &complex, &scalar,
721 )?));
722 }
723 (
724 TensorRead::View(TensorView::$complex_variant(complex)),
725 TensorRead::Tensor(Tensor::$real_variant(real)),
726 ) if real.shape().is_empty() => {
727 let scalar = complex_scalar_tensor_from_tensor(real)?;
728 let scalar = scalar.as_view();
729 return Ok(Tensor::$complex_variant(typed_sub_view_with_pool(
730 buffers, complex, &scalar,
731 )?));
732 }
733 (
734 TensorRead::View(TensorView::$complex_variant(complex)),
735 TensorRead::View(TensorView::$real_variant(real)),
736 ) if real.shape().is_empty() => {
737 let scalar = complex_scalar_tensor_from_view(real)?;
738 let scalar = scalar.as_view();
739 return Ok(Tensor::$complex_variant(typed_sub_view_with_pool(
740 buffers, complex, &scalar,
741 )?));
742 }
743 _ => {}
744 }
745 };
746 }
747
748 dispatch_real_complex_scalar!(F32, C32);
749 dispatch_real_complex_scalar!(F64, C64);
750
751 dispatch!(F32, typed_sub_view_with_pool);
752 dispatch!(F64, typed_sub_view_with_pool);
753 dispatch!(I32, typed_wrapping_sub_view_with_pool);
754 dispatch!(I64, typed_wrapping_sub_view_with_pool);
755 dispatch!(C32, typed_sub_view_with_pool);
756 dispatch!(C64, typed_sub_view_with_pool);
757
758 Err(read_pair_error("sub", lhs, rhs))
759}
760
761pub fn mul(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
776 with_local_pool(|buffers| mul_with_pool(buffers, lhs, rhs))
777}
778
779fn binary_read_with_pool(
780 op: &'static str,
781 buffers: &mut BufferPool,
782 lhs: TensorRead<'_>,
783 rhs: TensorRead<'_>,
784 f: impl FnOnce(&mut BufferPool, &Tensor, &Tensor) -> crate::Result<Tensor>,
785) -> crate::Result<Tensor> {
786 if let (Some(lhs), Some(rhs)) = (lhs.as_tensor(), rhs.as_tensor()) {
787 return f(buffers, lhs, rhs);
788 }
789
790 Err(read_pair_error(op, lhs, rhs))
791}
792
793pub(crate) fn mul_with_pool(
794 buffers: &mut BufferPool,
795 lhs: &Tensor,
796 rhs: &Tensor,
797) -> crate::Result<Tensor> {
798 match (lhs, rhs) {
799 (Tensor::F32(a), Tensor::F32(b)) => Ok(Tensor::F32(typed_mul_with_pool(buffers, a, b)?)),
800 (Tensor::F64(a), Tensor::F64(b)) => Ok(Tensor::F64(typed_mul_with_pool(buffers, a, b)?)),
801 (Tensor::I32(a), Tensor::I32(b)) => {
802 Ok(Tensor::I32(typed_wrapping_mul_with_pool(buffers, a, b)?))
803 }
804 (Tensor::I64(a), Tensor::I64(b)) => {
805 Ok(Tensor::I64(typed_wrapping_mul_with_pool(buffers, a, b)?))
806 }
807 (Tensor::C32(a), Tensor::C32(b)) => Ok(Tensor::C32(typed_mul_with_pool(buffers, a, b)?)),
808 (Tensor::C64(a), Tensor::C64(b)) => Ok(Tensor::C64(typed_mul_with_pool(buffers, a, b)?)),
809 (Tensor::F32(a), Tensor::C32(b)) if a.shape().is_empty() => {
810 let scalar = complex_scalar_tensor(typed_host_data("mul", a)?[0])?;
811 Ok(Tensor::C32(typed_mul_with_pool(buffers, &scalar, b)?))
812 }
813 (Tensor::C32(a), Tensor::F32(b)) if b.shape().is_empty() => {
814 let scalar = complex_scalar_tensor(typed_host_data("mul", b)?[0])?;
815 Ok(Tensor::C32(typed_mul_with_pool(buffers, a, &scalar)?))
816 }
817 (Tensor::F64(a), Tensor::C64(b)) if a.shape().is_empty() => {
818 let scalar = complex_scalar_tensor(typed_host_data("mul", a)?[0])?;
819 Ok(Tensor::C64(typed_mul_with_pool(buffers, &scalar, b)?))
820 }
821 (Tensor::C64(a), Tensor::F64(b)) if b.shape().is_empty() => {
822 let scalar = complex_scalar_tensor(typed_host_data("mul", b)?[0])?;
823 Ok(Tensor::C64(typed_mul_with_pool(buffers, a, &scalar)?))
824 }
825 _ => Err(tensor_pair_error("mul", lhs, rhs)),
826 }
827}
828
829pub(crate) fn mul_read_with_pool(
830 buffers: &mut BufferPool,
831 lhs: TensorRead<'_>,
832 rhs: TensorRead<'_>,
833) -> crate::Result<Tensor> {
834 if let (TensorRead::Tensor(lhs), TensorRead::Tensor(rhs)) = (&lhs, &rhs) {
835 return mul_with_pool(buffers, lhs, rhs);
836 }
837
838 macro_rules! dispatch {
839 ($variant:ident, $func:ident) => {
840 match (&lhs, &rhs) {
841 (
842 TensorRead::Tensor(Tensor::$variant(a)),
843 TensorRead::View(TensorView::$variant(b)),
844 ) => {
845 let a = a.as_view();
846 return Ok(Tensor::$variant($func(buffers, &a, b)?));
847 }
848 (
849 TensorRead::View(TensorView::$variant(a)),
850 TensorRead::Tensor(Tensor::$variant(b)),
851 ) => {
852 let b = b.as_view();
853 return Ok(Tensor::$variant($func(buffers, a, &b)?));
854 }
855 (
856 TensorRead::View(TensorView::$variant(a)),
857 TensorRead::View(TensorView::$variant(b)),
858 ) => {
859 return Ok(Tensor::$variant($func(buffers, a, b)?));
860 }
861 _ => {}
862 }
863 };
864 }
865
866 macro_rules! dispatch_real_complex_scalar {
867 ($real_variant:ident, $complex_variant:ident) => {
868 match (&lhs, &rhs) {
869 (
870 TensorRead::Tensor(Tensor::$real_variant(real)),
871 TensorRead::View(TensorView::$complex_variant(complex)),
872 ) if real.shape().is_empty() => {
873 let scalar = complex_scalar_tensor_from_tensor(real)?;
874 let scalar = scalar.as_view();
875 return Ok(Tensor::$complex_variant(typed_mul_view_with_pool(
876 buffers, &scalar, complex,
877 )?));
878 }
879 (
880 TensorRead::View(TensorView::$real_variant(real)),
881 TensorRead::Tensor(Tensor::$complex_variant(complex)),
882 ) if real.shape().is_empty() => {
883 let scalar = complex_scalar_tensor_from_view(real)?;
884 let scalar = scalar.as_view();
885 let complex = complex.as_view();
886 return Ok(Tensor::$complex_variant(typed_mul_view_with_pool(
887 buffers, &scalar, &complex,
888 )?));
889 }
890 (
891 TensorRead::View(TensorView::$real_variant(real)),
892 TensorRead::View(TensorView::$complex_variant(complex)),
893 ) if real.shape().is_empty() => {
894 let scalar = complex_scalar_tensor_from_view(real)?;
895 let scalar = scalar.as_view();
896 return Ok(Tensor::$complex_variant(typed_mul_view_with_pool(
897 buffers, &scalar, complex,
898 )?));
899 }
900 (
901 TensorRead::Tensor(Tensor::$complex_variant(complex)),
902 TensorRead::View(TensorView::$real_variant(real)),
903 ) if real.shape().is_empty() => {
904 let complex = complex.as_view();
905 let scalar = complex_scalar_tensor_from_view(real)?;
906 let scalar = scalar.as_view();
907 return Ok(Tensor::$complex_variant(typed_mul_view_with_pool(
908 buffers, &complex, &scalar,
909 )?));
910 }
911 (
912 TensorRead::View(TensorView::$complex_variant(complex)),
913 TensorRead::Tensor(Tensor::$real_variant(real)),
914 ) if real.shape().is_empty() => {
915 let scalar = complex_scalar_tensor_from_tensor(real)?;
916 let scalar = scalar.as_view();
917 return Ok(Tensor::$complex_variant(typed_mul_view_with_pool(
918 buffers, complex, &scalar,
919 )?));
920 }
921 (
922 TensorRead::View(TensorView::$complex_variant(complex)),
923 TensorRead::View(TensorView::$real_variant(real)),
924 ) if real.shape().is_empty() => {
925 let scalar = complex_scalar_tensor_from_view(real)?;
926 let scalar = scalar.as_view();
927 return Ok(Tensor::$complex_variant(typed_mul_view_with_pool(
928 buffers, complex, &scalar,
929 )?));
930 }
931 _ => {}
932 }
933 };
934 }
935
936 dispatch_real_complex_scalar!(F32, C32);
937 dispatch_real_complex_scalar!(F64, C64);
938
939 dispatch!(F32, typed_mul_view_with_pool);
940 dispatch!(F64, typed_mul_view_with_pool);
941 dispatch!(I32, typed_wrapping_mul_view_with_pool);
942 dispatch!(I64, typed_wrapping_mul_view_with_pool);
943 dispatch!(C32, typed_mul_view_with_pool);
944 dispatch!(C64, typed_mul_view_with_pool);
945
946 binary_read_with_pool("mul", buffers, lhs, rhs, mul_with_pool)
947}
948
949enum CpuReadView<'a> {
950 F32(TypedTensorView<'a, f32>),
951 F64(TypedTensorView<'a, f64>),
952 I32(TypedTensorView<'a, i32>),
953 I64(TypedTensorView<'a, i64>),
954 Bool(TypedTensorView<'a, bool>),
955 C32(TypedTensorView<'a, Complex<f32>>),
956 C64(TypedTensorView<'a, Complex<f64>>),
957}
958
959fn read_as_cpu_view(input: TensorRead<'_>) -> CpuReadView<'_> {
960 match input {
961 TensorRead::Tensor(Tensor::F32(tensor)) => CpuReadView::F32(tensor.as_view()),
962 TensorRead::Tensor(Tensor::F64(tensor)) => CpuReadView::F64(tensor.as_view()),
963 TensorRead::Tensor(Tensor::I32(tensor)) => CpuReadView::I32(tensor.as_view()),
964 TensorRead::Tensor(Tensor::I64(tensor)) => CpuReadView::I64(tensor.as_view()),
965 TensorRead::Tensor(Tensor::Bool(tensor)) => CpuReadView::Bool(tensor.as_view()),
966 TensorRead::Tensor(Tensor::C32(tensor)) => CpuReadView::C32(tensor.as_view()),
967 TensorRead::Tensor(Tensor::C64(tensor)) => CpuReadView::C64(tensor.as_view()),
968 TensorRead::View(TensorView::F32(view)) => CpuReadView::F32(view),
969 TensorRead::View(TensorView::F64(view)) => CpuReadView::F64(view),
970 TensorRead::View(TensorView::I32(view)) => CpuReadView::I32(view),
971 TensorRead::View(TensorView::I64(view)) => CpuReadView::I64(view),
972 TensorRead::View(TensorView::Bool(view)) => CpuReadView::Bool(view),
973 TensorRead::View(TensorView::C32(view)) => CpuReadView::C32(view),
974 TensorRead::View(TensorView::C64(view)) => CpuReadView::C64(view),
975 }
976}
977
978pub(crate) fn elementwise_fusion_with_pool(
979 buffers: &mut BufferPool,
980 inputs: &[&Tensor],
981 plan: &ElementwiseFusionPlan,
982) -> crate::Result<Option<Vec<Tensor>>> {
983 if !validate_elementwise_fusion_inputs(inputs, plan)? {
984 return Ok(None);
985 }
986 if inputs.is_empty() {
987 return Ok(None);
988 }
989 if plan_uses_unfused_op(plan) {
990 return Ok(None);
991 }
992
993 match plan.dtype() {
994 DType::F32 => {
995 let typed_inputs = inputs
996 .iter()
997 .map(|input| match input {
998 Tensor::F32(tensor) => Ok(tensor),
999 _ => Err(crate::Error::DTypeMismatch {
1000 op: ELEMENTWISE_FUSION_OP,
1001 lhs: input.dtype(),
1002 rhs: plan.dtype(),
1003 }),
1004 })
1005 .collect::<crate::Result<Vec<_>>>()?;
1006 typed_elementwise_fusion_with_pool(buffers, &typed_inputs, plan, Tensor::F32)
1007 }
1008 DType::F64 => {
1009 let typed_inputs = inputs
1010 .iter()
1011 .map(|input| match input {
1012 Tensor::F64(tensor) => Ok(tensor),
1013 _ => Err(crate::Error::DTypeMismatch {
1014 op: ELEMENTWISE_FUSION_OP,
1015 lhs: input.dtype(),
1016 rhs: plan.dtype(),
1017 }),
1018 })
1019 .collect::<crate::Result<Vec<_>>>()?;
1020 typed_elementwise_fusion_with_pool(buffers, &typed_inputs, plan, Tensor::F64)
1021 }
1022 DType::C32 => {
1023 if plan_uses_ordered_op(plan) {
1024 return Ok(None);
1025 }
1026 let typed_inputs = inputs
1027 .iter()
1028 .map(|input| match input {
1029 Tensor::C32(tensor) => Ok(tensor),
1030 _ => Err(crate::Error::DTypeMismatch {
1031 op: ELEMENTWISE_FUSION_OP,
1032 lhs: input.dtype(),
1033 rhs: plan.dtype(),
1034 }),
1035 })
1036 .collect::<crate::Result<Vec<_>>>()?;
1037 typed_elementwise_fusion_with_pool(buffers, &typed_inputs, plan, Tensor::C32)
1038 }
1039 DType::C64 => {
1040 if plan_uses_ordered_op(plan) {
1041 return Ok(None);
1042 }
1043 let typed_inputs = inputs
1044 .iter()
1045 .map(|input| match input {
1046 Tensor::C64(tensor) => Ok(tensor),
1047 _ => Err(crate::Error::DTypeMismatch {
1048 op: ELEMENTWISE_FUSION_OP,
1049 lhs: input.dtype(),
1050 rhs: plan.dtype(),
1051 }),
1052 })
1053 .collect::<crate::Result<Vec<_>>>()?;
1054 typed_elementwise_fusion_with_pool(buffers, &typed_inputs, plan, Tensor::C64)
1055 }
1056 DType::I32 | DType::I64 | DType::Bool => Ok(None),
1057 }
1058}
1059
1060fn typed_elementwise_fusion_with_pool<T>(
1061 buffers: &mut BufferPool,
1062 inputs: &[&TypedTensor<T>],
1063 plan: &ElementwiseFusionPlan,
1064 wrap: fn(TypedTensor<T>) -> Tensor,
1065) -> crate::Result<Option<Vec<Tensor>>>
1066where
1067 T: Copy + Clone + FusedScalar + PoolScalar,
1068{
1069 if should_defer_to_broadcast_multiply_special_case(plan) {
1070 return Ok(None);
1071 }
1072 if plan.input_views().iter().all(|view| view.is_identity()) {
1073 return typed_elementwise_fusion_identity_with_pool(buffers, inputs, plan, wrap);
1074 }
1075
1076 let input_views = inputs
1077 .iter()
1078 .zip(plan.input_views())
1079 .map(|(input, view)| typed_fusion_input_view(input, view))
1080 .collect::<crate::Result<Vec<_>>>()?;
1081 let shape = input_views[0].dims().to_vec();
1082 if input_views
1083 .iter()
1084 .skip(1)
1085 .any(|input| input.dims() != shape.as_slice())
1086 {
1087 return Ok(None);
1088 }
1089 let element_count =
1090 tenferro_tensor::validate::checked_shape_product(ELEMENTWISE_FUSION_OP, "shape", &shape)?;
1091 if element_count < ELEMENTWISE_FUSION_MIN_ELEMENTS {
1092 return Ok(None);
1093 }
1094
1095 run_typed_elementwise_fusion_with_views(buffers, &input_views, &shape, plan, wrap)
1096}
1097
1098fn typed_elementwise_fusion_identity_with_pool<T>(
1099 buffers: &mut BufferPool,
1100 inputs: &[&TypedTensor<T>],
1101 plan: &ElementwiseFusionPlan,
1102 wrap: fn(TypedTensor<T>) -> Tensor,
1103) -> crate::Result<Option<Vec<Tensor>>>
1104where
1105 T: Copy + Clone + FusedScalar + PoolScalar,
1106{
1107 let shape = inputs[0].shape();
1108 if inputs.iter().skip(1).any(|input| input.shape() != shape) {
1109 return Ok(None);
1110 }
1111 let element_count =
1112 tenferro_tensor::validate::checked_shape_product(ELEMENTWISE_FUSION_OP, "shape", shape)?;
1113 if element_count < ELEMENTWISE_FUSION_MIN_ELEMENTS {
1114 return Ok(None);
1115 }
1116
1117 let input_views = inputs
1118 .iter()
1119 .map(|input| typed_view(ELEMENTWISE_FUSION_OP, input))
1120 .collect::<crate::Result<Vec<_>>>()?;
1121 run_typed_elementwise_fusion_with_views(buffers, &input_views, shape, plan, wrap)
1122}
1123
1124fn run_typed_elementwise_fusion_with_views<T>(
1125 buffers: &mut BufferPool,
1126 input_views: &[StridedView<'_, T>],
1127 shape: &[usize],
1128 plan: &ElementwiseFusionPlan,
1129 wrap: fn(TypedTensor<T>) -> Tensor,
1130) -> crate::Result<Option<Vec<Tensor>>>
1131where
1132 T: Copy + Clone + FusedScalar + PoolScalar,
1133{
1134 if let Some(outputs) =
1135 try_typed_mul_add_specialization(buffers, input_views, shape, plan, wrap)?
1136 {
1137 return Ok(Some(outputs));
1138 }
1139
1140 let fused_plan = strided_fused_plan(plan);
1141 let mut output_arrays = Vec::with_capacity(plan.outputs().len());
1142 for _ in plan.outputs() {
1143 output_arrays.push(unsafe { typed_array_uninit_from_pool(buffers, shape) }?);
1145 }
1146
1147 {
1148 let mut output_views = output_arrays
1149 .iter_mut()
1150 .map(|output| output.view_mut())
1151 .collect::<Vec<_>>();
1152 fused_elementwise_into(&mut output_views, input_views, &fused_plan)
1153 .map_err(|err| crate::Error::backend_failure(ELEMENTWISE_FUSION_OP, err))?;
1154 }
1155
1156 Ok(Some(
1157 output_arrays
1158 .into_iter()
1159 .map(|output| wrap(tensor_from_array(output)))
1160 .collect(),
1161 ))
1162}
1163
1164fn try_typed_mul_add_specialization<T>(
1165 buffers: &mut BufferPool,
1166 input_views: &[StridedView<'_, T>],
1167 shape: &[usize],
1168 plan: &ElementwiseFusionPlan,
1169 wrap: fn(TypedTensor<T>) -> Tensor,
1170) -> crate::Result<Option<Vec<Tensor>>>
1171where
1172 T: Copy + Clone + FusedScalar + PoolScalar,
1173{
1174 if plan.input_count() != 2
1175 || input_views.len() != 2
1176 || plan.outputs() != [3]
1177 || plan.ops().len() != 2
1178 || plan.ops()[0].op() != ElementwiseFusionOp::Multiply
1179 || plan.ops()[0].inputs() != [0, 1]
1180 || plan.ops()[1].op() != ElementwiseFusionOp::Add
1181 {
1182 return Ok(None);
1183 }
1184
1185 let mut out = unsafe { typed_array_uninit_from_pool(buffers, shape) }?;
1187 match plan.ops()[1].inputs() {
1188 [2, 0] | [0, 2] => zip_map2_into(
1189 &mut out.view_mut(),
1190 &input_views[0],
1191 &input_views[1],
1192 |a, b| a.fused_multiply(b).fused_add(a),
1193 ),
1194 [2, 1] | [1, 2] => zip_map2_into(
1195 &mut out.view_mut(),
1196 &input_views[0],
1197 &input_views[1],
1198 |a, b| a.fused_multiply(b).fused_add(b),
1199 ),
1200 _ => return Ok(None),
1201 }
1202 .map_err(|err| crate::Error::backend_failure(ELEMENTWISE_FUSION_OP, err))?;
1203
1204 Ok(Some(vec![wrap(tensor_from_array(out))]))
1205}
1206
1207fn typed_fusion_input_view<'a, T>(
1208 input: &'a TypedTensor<T>,
1209 view: &ElementwiseFusionInputView,
1210) -> crate::Result<StridedView<'a, T>>
1211where
1212 T: Copy,
1213{
1214 let base = typed_view(ELEMENTWISE_FUSION_OP, input)?;
1215 let ElementwiseFusionInputView::BroadcastInDim { shape, dims } = view else {
1216 return Ok(base);
1217 };
1218 if dims.len() != input.shape().len() {
1219 return Err(crate::Error::InvalidConfig {
1220 op: ELEMENTWISE_FUSION_OP,
1221 message: format!(
1222 "broadcast dims length {} does not match input rank {}",
1223 dims.len(),
1224 input.shape().len()
1225 ),
1226 });
1227 }
1228
1229 let mut strides = vec![0; shape.len()];
1230 let mut seen = Vec::with_capacity(shape.len());
1231 seen.resize(shape.len(), false);
1232 for (source_axis, &target_axis) in dims.iter().enumerate() {
1233 if target_axis >= shape.len() {
1234 return Err(crate::Error::AxisOutOfBounds {
1235 op: ELEMENTWISE_FUSION_OP,
1236 axis: target_axis,
1237 rank: shape.len(),
1238 });
1239 }
1240 if seen[target_axis] {
1241 return Err(crate::Error::DuplicateAxis {
1242 op: ELEMENTWISE_FUSION_OP,
1243 axis: target_axis,
1244 role: "broadcast dims",
1245 });
1246 }
1247 seen[target_axis] = true;
1248 let source_dim = input.shape()[source_axis];
1249 let target_dim = shape[target_axis];
1250 if source_dim != target_dim && source_dim != 1 {
1251 return Err(crate::Error::ShapeMismatch {
1252 op: ELEMENTWISE_FUSION_OP,
1253 lhs: shape.to_vec(),
1254 rhs: input.shape().to_vec(),
1255 });
1256 }
1257 if source_dim == target_dim {
1258 strides[target_axis] = base.strides()[source_axis];
1259 }
1260 }
1261
1262 StridedView::new(base.data(), shape, &strides, base.offset())
1263 .map_err(|err| crate::Error::backend_failure(ELEMENTWISE_FUSION_OP, err))
1264}
1265
1266fn typed_binary_view_with_pool<T, L, R>(
1267 op: &'static str,
1268 buffers: &mut BufferPool,
1269 lhs: &TypedTensorView<'_, T, L>,
1270 rhs: &TypedTensorView<'_, T, R>,
1271 f: impl Fn(T, T) -> T + Copy + Sync,
1272) -> crate::Result<TypedTensor<T>>
1273where
1274 T: Copy + PoolScalar + 'static,
1275 L: TensorRank,
1276 R: TensorRank,
1277{
1278 if lhs.shape() == rhs.shape() {
1279 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
1281 zip_map2_into(
1282 &mut out.view_mut(),
1283 &typed_view_from_view(op, lhs)?,
1284 &typed_view_from_view(op, rhs)?,
1285 f,
1286 )
1287 .map_err(|err| crate::Error::backend_failure(op, err))?;
1288 Ok(tensor_from_array(out))
1289 } else if lhs.shape().is_empty() {
1290 let scalar = typed_view_from_view(op, lhs)?.get(&[]);
1291 let mut out = unsafe { typed_array_uninit_from_pool(buffers, rhs.shape()) }?;
1293 map_into(&mut out.view_mut(), &typed_view_from_view(op, rhs)?, |x| {
1294 f(scalar, x)
1295 })
1296 .map_err(|err| crate::Error::backend_failure(op, err))?;
1297 Ok(tensor_from_array(out))
1298 } else if rhs.shape().is_empty() {
1299 let scalar = typed_view_from_view(op, rhs)?.get(&[]);
1300 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
1302 map_into(&mut out.view_mut(), &typed_view_from_view(op, lhs)?, |x| {
1303 f(x, scalar)
1304 })
1305 .map_err(|err| crate::Error::backend_failure(op, err))?;
1306 Ok(tensor_from_array(out))
1307 } else {
1308 Err(crate::Error::ShapeMismatch {
1309 op,
1310 lhs: lhs.shape().to_vec(),
1311 rhs: rhs.shape().to_vec(),
1312 })
1313 }
1314}
1315
1316fn typed_unary_view_with_pool<T, R>(
1317 op: &'static str,
1318 buffers: &mut BufferPool,
1319 input: &TypedTensorView<'_, T, R>,
1320 f: impl Fn(T) -> T + Copy + Sync,
1321) -> crate::Result<TypedTensor<T>>
1322where
1323 T: Copy + PoolScalar + 'static,
1324 R: TensorRank,
1325{
1326 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
1328 map_into(&mut out.view_mut(), &typed_view_from_view(op, input)?, f)
1329 .map_err(|err| crate::Error::backend_failure(op, err))?;
1330 Ok(tensor_from_array(out))
1331}
1332
1333fn typed_same_shape_binary_view_with_pool<T, O, L, R>(
1334 op: &'static str,
1335 buffers: &mut BufferPool,
1336 lhs: &TypedTensorView<'_, T, L>,
1337 rhs: &TypedTensorView<'_, T, R>,
1338 f: impl Fn(T, T) -> O + Copy + Sync,
1339) -> crate::Result<TypedTensor<O>>
1340where
1341 T: Copy + Send + Sync + 'static,
1342 O: Copy + PoolScalar,
1343 L: TensorRank,
1344 R: TensorRank,
1345{
1346 if lhs.shape() != rhs.shape() {
1347 return Err(crate::Error::ShapeMismatch {
1348 op,
1349 lhs: lhs.shape().to_vec(),
1350 rhs: rhs.shape().to_vec(),
1351 });
1352 }
1353 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
1355 zip_map2_into(
1356 &mut out.view_mut(),
1357 &typed_view_from_view(op, lhs)?,
1358 &typed_view_from_view(op, rhs)?,
1359 f,
1360 )
1361 .map_err(|err| crate::Error::backend_failure(op, err))?;
1362 Ok(tensor_from_array(out))
1363}
1364
1365fn typed_select_view_with_pool<T, P, A, B>(
1366 buffers: &mut BufferPool,
1367 pred: &TypedTensorView<'_, bool, P>,
1368 on_true: &TypedTensorView<'_, T, A>,
1369 on_false: &TypedTensorView<'_, T, B>,
1370) -> crate::Result<TypedTensor<T>>
1371where
1372 T: Copy + PoolScalar + 'static,
1373 P: TensorRank,
1374 A: TensorRank,
1375 B: TensorRank,
1376{
1377 if pred.shape() != on_true.shape() {
1378 return Err(crate::Error::ShapeMismatch {
1379 op: "select",
1380 lhs: pred.shape().to_vec(),
1381 rhs: on_true.shape().to_vec(),
1382 });
1383 }
1384 if pred.shape() != on_false.shape() {
1385 return Err(crate::Error::ShapeMismatch {
1386 op: "select",
1387 lhs: pred.shape().to_vec(),
1388 rhs: on_false.shape().to_vec(),
1389 });
1390 }
1391 let mut out = unsafe { typed_array_uninit_from_pool(buffers, pred.shape()) }?;
1393 zip_map3_into(
1394 &mut out.view_mut(),
1395 &typed_view_from_view("select", pred)?,
1396 &typed_view_from_view("select", on_true)?,
1397 &typed_view_from_view("select", on_false)?,
1398 |p, t, f| if p { t } else { f },
1399 )
1400 .map_err(|err| crate::Error::backend_failure("select", err))?;
1401 Ok(tensor_from_array(out))
1402}
1403
1404fn typed_clamp_view_with_pool<T, I, L, U>(
1405 buffers: &mut BufferPool,
1406 input: &TypedTensorView<'_, T, I>,
1407 lower: &TypedTensorView<'_, T, L>,
1408 upper: &TypedTensorView<'_, T, U>,
1409) -> crate::Result<TypedTensor<T>>
1410where
1411 T: OrderedElem + PoolScalar + 'static,
1412 I: TensorRank,
1413 L: TensorRank,
1414 U: TensorRank,
1415{
1416 if input.shape() != lower.shape() {
1417 return Err(crate::Error::ShapeMismatch {
1418 op: "clamp",
1419 lhs: input.shape().to_vec(),
1420 rhs: lower.shape().to_vec(),
1421 });
1422 }
1423 if input.shape() != upper.shape() {
1424 return Err(crate::Error::ShapeMismatch {
1425 op: "clamp",
1426 lhs: input.shape().to_vec(),
1427 rhs: upper.shape().to_vec(),
1428 });
1429 }
1430 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
1432 zip_map3_into(
1433 &mut out.view_mut(),
1434 &typed_view_from_view("clamp", input)?,
1435 &typed_view_from_view("clamp", lower)?,
1436 &typed_view_from_view("clamp", upper)?,
1437 |x, lo, hi| hi.min_elem(lo.max_elem(x)),
1438 )
1439 .map_err(|err| crate::Error::backend_failure("clamp", err))?;
1440 Ok(tensor_from_array(out))
1441}
1442
1443#[derive(Clone, Copy)]
1444enum SplitOuterProductLayout {
1445 LhsPrefix,
1446 RhsPrefix,
1447}
1448
1449struct SplitOuterProductPlan {
1450 #[allow(dead_code)]
1451 rows: usize,
1452 #[allow(dead_code)]
1453 cols: usize,
1454 #[allow(dead_code)]
1455 batches: usize,
1456 layout: SplitOuterProductLayout,
1457 lhs_free_axes: Vec<usize>,
1458 rhs_free_axes: Vec<usize>,
1459 lhs_batch_axes: Vec<usize>,
1460 rhs_batch_axes: Vec<usize>,
1461}
1462
1463struct OuterProductAxisPartition {
1464 lhs_free_output_axes: Vec<usize>,
1465 rhs_free_output_axes: Vec<usize>,
1466 batch_output_axes: Vec<usize>,
1467 lhs_free_axes: Vec<usize>,
1468 rhs_free_axes: Vec<usize>,
1469 lhs_batch_axes: Vec<usize>,
1470 rhs_batch_axes: Vec<usize>,
1471}
1472
1473fn shape_matches_dims(source_shape: &[usize], output_shape: &[usize], dims: &[usize]) -> bool {
1474 source_shape.len() == dims.len()
1475 && source_shape
1476 .iter()
1477 .zip(dims.iter())
1478 .all(|(&dim, &axis)| output_shape.get(axis).copied() == Some(dim))
1479}
1480
1481fn axes_by_output(dims: &[usize], output_rank: usize) -> Option<Vec<Option<usize>>> {
1482 let mut axes = vec![None; output_rank];
1483 for (src_axis, &dst_axis) in dims.iter().enumerate() {
1484 let slot = axes.get_mut(dst_axis)?;
1485 if slot.replace(src_axis).is_some() {
1486 return None;
1487 }
1488 }
1489 Some(axes)
1490}
1491
1492fn axes_shape_product<T>(
1493 op: &'static str,
1494 view: &TypedTensorView<'_, T>,
1495 axes: &[usize],
1496) -> crate::Result<usize>
1497where
1498 T: 'static,
1499{
1500 axes.iter().try_fold(1usize, |acc, &axis| {
1501 acc.checked_mul(view.shape()[axis])
1502 .ok_or_else(|| crate::Error::backend_failure(op, "shape size overflows usize"))
1503 })
1504}
1505
1506fn classify_outer_product_axes(
1507 lhs_dims: &[usize],
1508 rhs_dims: &[usize],
1509 output_rank: usize,
1510) -> Option<OuterProductAxisPartition> {
1511 let lhs_axes_by_output = axes_by_output(lhs_dims, output_rank)?;
1512 let rhs_axes_by_output = axes_by_output(rhs_dims, output_rank)?;
1513
1514 let mut lhs_free_output_axes = Vec::new();
1515 let mut rhs_free_output_axes = Vec::new();
1516 let mut batch_output_axes = Vec::new();
1517 let mut lhs_free_axes = Vec::new();
1518 let mut rhs_free_axes = Vec::new();
1519 let mut lhs_batch_axes = Vec::new();
1520 let mut rhs_batch_axes = Vec::new();
1521
1522 for output_axis in 0..output_rank {
1523 match (
1524 lhs_axes_by_output[output_axis],
1525 rhs_axes_by_output[output_axis],
1526 ) {
1527 (Some(lhs_axis), Some(rhs_axis)) => {
1528 batch_output_axes.push(output_axis);
1529 lhs_batch_axes.push(lhs_axis);
1530 rhs_batch_axes.push(rhs_axis);
1531 }
1532 (Some(lhs_axis), None) => {
1533 lhs_free_output_axes.push(output_axis);
1534 lhs_free_axes.push(lhs_axis);
1535 }
1536 (None, Some(rhs_axis)) => {
1537 rhs_free_output_axes.push(output_axis);
1538 rhs_free_axes.push(rhs_axis);
1539 }
1540 (None, None) => return None,
1541 }
1542 }
1543
1544 Some(OuterProductAxisPartition {
1545 lhs_free_output_axes,
1546 rhs_free_output_axes,
1547 batch_output_axes,
1548 lhs_free_axes,
1549 rhs_free_axes,
1550 lhs_batch_axes,
1551 rhs_batch_axes,
1552 })
1553}
1554
1555fn output_axes_match_partition(output_rank: usize, groups: &[&[usize]]) -> bool {
1556 groups
1557 .iter()
1558 .flat_map(|group| group.iter().copied())
1559 .eq(0..output_rank)
1560}
1561
1562fn split_outer_product_plan<T>(
1563 lhs: &TypedTensorView<'_, T>,
1564 lhs_shape: &[usize],
1565 lhs_dims: &[usize],
1566 rhs: &TypedTensorView<'_, T>,
1567 rhs_shape: &[usize],
1568 rhs_dims: &[usize],
1569) -> crate::Result<Option<SplitOuterProductPlan>>
1570where
1571 T: 'static,
1572{
1573 let output_rank = lhs_shape.len();
1574 if lhs_shape != rhs_shape
1575 || !shape_matches_dims(lhs.shape(), lhs_shape, lhs_dims)
1576 || !shape_matches_dims(rhs.shape(), rhs_shape, rhs_dims)
1577 || lhs.backend_buffer().is_some()
1578 || rhs.backend_buffer().is_some()
1579 || lhs.offset() < 0
1580 || rhs.offset() < 0
1581 || lhs.strides().iter().any(|&stride| stride < 0)
1582 || rhs.strides().iter().any(|&stride| stride < 0)
1583 {
1584 return Ok(None);
1585 }
1586
1587 let Some(partition) = classify_outer_product_axes(lhs_dims, rhs_dims, output_rank) else {
1588 return Ok(None);
1589 };
1590
1591 let lhs_free_size = axes_shape_product("broadcast_multiply", lhs, &partition.lhs_free_axes)?;
1592 let rhs_free_size = axes_shape_product("broadcast_multiply", rhs, &partition.rhs_free_axes)?;
1593 if lhs_free_size <= 1 || rhs_free_size <= 1 {
1594 return Ok(None);
1595 }
1596 let batches = axes_shape_product("broadcast_multiply", lhs, &partition.lhs_batch_axes)?;
1597
1598 let lhs_prefix = output_axes_match_partition(
1599 output_rank,
1600 &[
1601 &partition.lhs_free_output_axes,
1602 &partition.rhs_free_output_axes,
1603 &partition.batch_output_axes,
1604 ],
1605 );
1606 if lhs_prefix {
1607 return Ok(Some(SplitOuterProductPlan {
1608 rows: lhs_free_size,
1609 cols: rhs_free_size,
1610 batches,
1611 layout: SplitOuterProductLayout::LhsPrefix,
1612 lhs_free_axes: partition.lhs_free_axes,
1613 rhs_free_axes: partition.rhs_free_axes,
1614 lhs_batch_axes: partition.lhs_batch_axes,
1615 rhs_batch_axes: partition.rhs_batch_axes,
1616 }));
1617 }
1618
1619 let rhs_prefix = output_axes_match_partition(
1620 output_rank,
1621 &[
1622 &partition.rhs_free_output_axes,
1623 &partition.lhs_free_output_axes,
1624 &partition.batch_output_axes,
1625 ],
1626 );
1627 if rhs_prefix {
1628 return Ok(Some(SplitOuterProductPlan {
1629 rows: rhs_free_size,
1630 cols: lhs_free_size,
1631 batches,
1632 layout: SplitOuterProductLayout::RhsPrefix,
1633 lhs_free_axes: partition.lhs_free_axes,
1634 rhs_free_axes: partition.rhs_free_axes,
1635 lhs_batch_axes: partition.lhs_batch_axes,
1636 rhs_batch_axes: partition.rhs_batch_axes,
1637 }));
1638 }
1639
1640 Ok(None)
1641}
1642
1643fn try_outer_product_with_pool<T>(
1644 buffers: &mut BufferPool,
1645 lhs: &TypedTensorView<'_, T>,
1646 lhs_shape: &[usize],
1647 lhs_dims: &[usize],
1648 rhs: &TypedTensorView<'_, T>,
1649 rhs_shape: &[usize],
1650 rhs_dims: &[usize],
1651) -> crate::Result<Option<TypedTensor<T>>>
1652where
1653 T: Copy + Clone + Mul<Output = T> + PoolScalar + 'static,
1654{
1655 let Some(plan) = split_outer_product_plan(lhs, lhs_shape, lhs_dims, rhs, rhs_shape, rhs_dims)?
1656 else {
1657 return Ok(None);
1658 };
1659
1660 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs_shape) }?;
1662 let lhs_view = typed_view_from_view("broadcast_multiply", lhs)?;
1663 let rhs_view = typed_view_from_view("broadcast_multiply", rhs)?;
1664 match plan.layout {
1665 SplitOuterProductLayout::LhsPrefix => {
1666 let lhs_perm: Vec<_> = plan
1667 .lhs_free_axes
1668 .iter()
1669 .chain(plan.lhs_batch_axes.iter())
1670 .copied()
1671 .collect();
1672 let rhs_perm: Vec<_> = plan
1673 .rhs_free_axes
1674 .iter()
1675 .chain(plan.rhs_batch_axes.iter())
1676 .copied()
1677 .collect();
1678 let lhs_outer = lhs_view
1679 .permute(&lhs_perm)
1680 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1681 let rhs_outer = rhs_view
1682 .permute(&rhs_perm)
1683 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1684 batched_outer_product_into(
1685 &mut out.view_mut(),
1686 &lhs_outer,
1687 &rhs_outer,
1688 plan.lhs_free_axes.len(),
1689 plan.rhs_free_axes.len(),
1690 )
1691 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1692 }
1693 SplitOuterProductLayout::RhsPrefix => {
1694 let lhs_perm: Vec<_> = plan
1695 .lhs_free_axes
1696 .iter()
1697 .chain(plan.lhs_batch_axes.iter())
1698 .copied()
1699 .collect();
1700 let rhs_perm: Vec<_> = plan
1701 .rhs_free_axes
1702 .iter()
1703 .chain(plan.rhs_batch_axes.iter())
1704 .copied()
1705 .collect();
1706 let lhs_outer = lhs_view
1707 .permute(&lhs_perm)
1708 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1709 let rhs_outer = rhs_view
1710 .permute(&rhs_perm)
1711 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1712 batched_outer_product_into(
1713 &mut out.view_mut(),
1714 &rhs_outer,
1715 &lhs_outer,
1716 plan.rhs_free_axes.len(),
1717 plan.lhs_free_axes.len(),
1718 )
1719 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1720 }
1721 }
1722 Ok(Some(tensor_from_array(out)))
1723}
1724
1725struct LazyOuterProduct<T> {
1726 base: TypedTensor<T>,
1727 shape: Vec<usize>,
1728 strides: Vec<isize>,
1729}
1730
1731fn axes_by_physical_stride<T>(view: &TypedTensorView<'_, T>, axes: &[usize]) -> Vec<usize>
1732where
1733 T: 'static,
1734{
1735 let mut sorted = axes.to_vec();
1736 sorted.sort_by(|&lhs_axis, &rhs_axis| {
1737 view.strides()[lhs_axis]
1738 .cmp(&view.strides()[rhs_axis])
1739 .then_with(|| lhs_axis.cmp(&rhs_axis))
1740 });
1741 sorted
1742}
1743
1744fn append_axis_shapes<T>(shape: &mut Vec<usize>, view: &TypedTensorView<'_, T>, axes: &[usize])
1745where
1746 T: 'static,
1747{
1748 shape.extend(axes.iter().map(|&axis| view.shape()[axis]));
1749}
1750
1751fn set_lazy_stride(
1752 logical_strides: &mut [Option<isize>],
1753 output_axis: usize,
1754 stride: isize,
1755) -> crate::Result<()> {
1756 let rank = logical_strides.len();
1757 let slot = logical_strides
1758 .get_mut(output_axis)
1759 .ok_or(crate::Error::AxisOutOfBounds {
1760 op: "broadcast_multiply",
1761 axis: output_axis,
1762 rank,
1763 })?;
1764 if slot.replace(stride).is_some() {
1765 return Err(crate::Error::DuplicateAxis {
1766 op: "broadcast_multiply",
1767 axis: output_axis,
1768 role: "lazy output layout",
1769 });
1770 }
1771 Ok(())
1772}
1773
1774struct LazyOuterProductStrideSpec<'a> {
1775 output_shape: &'a [usize],
1776 base_shape: &'a [usize],
1777 leading_axes: &'a [usize],
1778 leading_dims: &'a [usize],
1779 trailing_axes: &'a [usize],
1780 trailing_dims: &'a [usize],
1781 lhs_batch_axes: &'a [usize],
1782 rhs_batch_axes: &'a [usize],
1783 lhs_dims: &'a [usize],
1784 rhs_dims: &'a [usize],
1785}
1786
1787fn lazy_outer_product_strides(spec: LazyOuterProductStrideSpec<'_>) -> crate::Result<Vec<isize>> {
1788 let base_strides = col_major_strides(spec.base_shape)?;
1789 let mut logical_strides = vec![None; spec.output_shape.len()];
1790 let mut base_axis = 0usize;
1791
1792 for &axis in spec.leading_axes {
1793 set_lazy_stride(
1794 &mut logical_strides,
1795 spec.leading_dims[axis],
1796 base_strides[base_axis],
1797 )?;
1798 base_axis += 1;
1799 }
1800 for &axis in spec.trailing_axes {
1801 set_lazy_stride(
1802 &mut logical_strides,
1803 spec.trailing_dims[axis],
1804 base_strides[base_axis],
1805 )?;
1806 base_axis += 1;
1807 }
1808 for (&lhs_axis, &rhs_axis) in spec.lhs_batch_axes.iter().zip(spec.rhs_batch_axes.iter()) {
1809 let output_axis = spec.lhs_dims[lhs_axis];
1810 if spec.rhs_dims[rhs_axis] != output_axis {
1811 return Err(crate::Error::backend_failure(
1812 "broadcast_multiply",
1813 "batch axes disagree while building lazy outer-product layout",
1814 ));
1815 }
1816 set_lazy_stride(&mut logical_strides, output_axis, base_strides[base_axis])?;
1817 base_axis += 1;
1818 }
1819
1820 logical_strides
1821 .into_iter()
1822 .collect::<Option<Vec<_>>>()
1823 .ok_or_else(|| {
1824 crate::Error::backend_failure(
1825 "broadcast_multiply",
1826 "lazy outer-product layout did not cover every output axis",
1827 )
1828 })
1829}
1830
1831fn lazy_outer_product_value(
1832 tensor: Tensor,
1833 shape: Vec<usize>,
1834 strides: Vec<isize>,
1835) -> crate::Result<TensorValue> {
1836 Ok(TensorValue::View(TensorOwnedView::from_parts(
1837 Arc::new(tensor),
1838 shape,
1839 strides,
1840 0,
1841 )?))
1842}
1843
1844fn try_lazy_outer_product_with_pool<T>(
1845 buffers: &mut BufferPool,
1846 lhs: &TypedTensorView<'_, T>,
1847 lhs_shape: &[usize],
1848 lhs_dims: &[usize],
1849 rhs: &TypedTensorView<'_, T>,
1850 rhs_shape: &[usize],
1851 rhs_dims: &[usize],
1852) -> crate::Result<Option<LazyOuterProduct<T>>>
1853where
1854 T: Copy + Clone + Mul<Output = T> + PoolScalar + 'static,
1855{
1856 let Some(plan) = split_outer_product_plan(lhs, lhs_shape, lhs_dims, rhs, rhs_shape, rhs_dims)?
1857 else {
1858 return Ok(None);
1859 };
1860
1861 let lhs_free_axes = axes_by_physical_stride(lhs, &plan.lhs_free_axes);
1862 let rhs_free_axes = axes_by_physical_stride(rhs, &plan.rhs_free_axes);
1863 if lhs_free_axes == plan.lhs_free_axes && rhs_free_axes == plan.rhs_free_axes {
1864 return Ok(None);
1865 }
1866
1867 let lhs_view = typed_view_from_view("broadcast_multiply", lhs)?;
1868 let rhs_view = typed_view_from_view("broadcast_multiply", rhs)?;
1869
1870 match plan.layout {
1871 SplitOuterProductLayout::LhsPrefix => {
1872 let lhs_perm: Vec<_> = lhs_free_axes
1873 .iter()
1874 .chain(plan.lhs_batch_axes.iter())
1875 .copied()
1876 .collect();
1877 let rhs_perm: Vec<_> = rhs_free_axes
1878 .iter()
1879 .chain(plan.rhs_batch_axes.iter())
1880 .copied()
1881 .collect();
1882 let lhs_outer = lhs_view
1883 .permute(&lhs_perm)
1884 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1885 let rhs_outer = rhs_view
1886 .permute(&rhs_perm)
1887 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1888
1889 let mut base_shape = Vec::with_capacity(lhs_shape.len());
1890 append_axis_shapes(&mut base_shape, lhs, &lhs_free_axes);
1891 append_axis_shapes(&mut base_shape, rhs, &rhs_free_axes);
1892 append_axis_shapes(&mut base_shape, lhs, &plan.lhs_batch_axes);
1893 let strides = lazy_outer_product_strides(LazyOuterProductStrideSpec {
1894 output_shape: lhs_shape,
1895 base_shape: &base_shape,
1896 leading_axes: &lhs_free_axes,
1897 leading_dims: lhs_dims,
1898 trailing_axes: &rhs_free_axes,
1899 trailing_dims: rhs_dims,
1900 lhs_batch_axes: &plan.lhs_batch_axes,
1901 rhs_batch_axes: &plan.rhs_batch_axes,
1902 lhs_dims,
1903 rhs_dims,
1904 })?;
1905
1906 let mut base = unsafe { typed_array_uninit_from_pool(buffers, &base_shape) }?;
1908 batched_outer_product_into(
1909 &mut base.view_mut(),
1910 &lhs_outer,
1911 &rhs_outer,
1912 lhs_free_axes.len(),
1913 rhs_free_axes.len(),
1914 )
1915 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1916 Ok(Some(LazyOuterProduct {
1917 base: tensor_from_array(base),
1918 shape: lhs_shape.to_vec(),
1919 strides,
1920 }))
1921 }
1922 SplitOuterProductLayout::RhsPrefix => {
1923 let lhs_perm: Vec<_> = lhs_free_axes
1924 .iter()
1925 .chain(plan.lhs_batch_axes.iter())
1926 .copied()
1927 .collect();
1928 let rhs_perm: Vec<_> = rhs_free_axes
1929 .iter()
1930 .chain(plan.rhs_batch_axes.iter())
1931 .copied()
1932 .collect();
1933 let lhs_outer = lhs_view
1934 .permute(&lhs_perm)
1935 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1936 let rhs_outer = rhs_view
1937 .permute(&rhs_perm)
1938 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1939
1940 let mut base_shape = Vec::with_capacity(lhs_shape.len());
1941 append_axis_shapes(&mut base_shape, rhs, &rhs_free_axes);
1942 append_axis_shapes(&mut base_shape, lhs, &lhs_free_axes);
1943 append_axis_shapes(&mut base_shape, lhs, &plan.lhs_batch_axes);
1944 let strides = lazy_outer_product_strides(LazyOuterProductStrideSpec {
1945 output_shape: lhs_shape,
1946 base_shape: &base_shape,
1947 leading_axes: &rhs_free_axes,
1948 leading_dims: rhs_dims,
1949 trailing_axes: &lhs_free_axes,
1950 trailing_dims: lhs_dims,
1951 lhs_batch_axes: &plan.lhs_batch_axes,
1952 rhs_batch_axes: &plan.rhs_batch_axes,
1953 lhs_dims,
1954 rhs_dims,
1955 })?;
1956
1957 let mut base = unsafe { typed_array_uninit_from_pool(buffers, &base_shape) }?;
1959 batched_outer_product_into(
1960 &mut base.view_mut(),
1961 &rhs_outer,
1962 &lhs_outer,
1963 rhs_free_axes.len(),
1964 lhs_free_axes.len(),
1965 )
1966 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
1967 Ok(Some(LazyOuterProduct {
1968 base: tensor_from_array(base),
1969 shape: lhs_shape.to_vec(),
1970 strides,
1971 }))
1972 }
1973 }
1974}
1975
1976#[allow(clippy::too_many_arguments)]
1977fn typed_broadcast_mul_view_with_pool<T, L, R>(
1978 buffers: &mut BufferPool,
1979 lhs: &TypedTensorView<'_, T, L>,
1980 lhs_shape: &[usize],
1981 lhs_dims: &[usize],
1982 rhs: &TypedTensorView<'_, T, R>,
1983 rhs_shape: &[usize],
1984 rhs_dims: &[usize],
1985) -> crate::Result<TypedTensor<T>>
1986where
1987 T: Copy + Clone + Zero + Mul<Output = T> + PoolScalar + 'static,
1988 L: TensorRank,
1989 R: TensorRank,
1990{
1991 if lhs_shape != rhs_shape {
1992 return Err(crate::Error::ShapeMismatch {
1993 op: "broadcast_multiply",
1994 lhs: lhs_shape.to_vec(),
1995 rhs: rhs_shape.to_vec(),
1996 });
1997 }
1998 let output_rank = lhs_shape.len();
1999 let lhs_is_scalar = lhs.shape().is_empty() && lhs_dims.is_empty();
2000 let rhs_is_scalar = rhs.shape().is_empty() && rhs_dims.is_empty();
2001 let lhs_is_full_output =
2002 lhs.shape() == lhs_shape && lhs_dims.iter().copied().eq(0..output_rank);
2003 let rhs_is_full_output =
2004 rhs.shape() == rhs_shape && rhs_dims.iter().copied().eq(0..output_rank);
2005 if lhs_is_scalar && rhs_is_scalar {
2006 let lhs_scalar = typed_view_from_view("broadcast_multiply", lhs)?.get(&[]);
2007 let rhs_scalar = typed_view_from_view("broadcast_multiply", rhs)?.get(&[]);
2008 return filled_broadcast_multiply_tensor(buffers, lhs_shape, lhs_scalar * rhs_scalar);
2009 }
2010 if lhs_is_scalar && rhs_is_full_output {
2011 let scalar = typed_view_from_view("broadcast_multiply", lhs)?.get(&[]);
2012 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs_shape) }?;
2014 map_into(
2015 &mut out.view_mut(),
2016 &typed_view_from_view("broadcast_multiply", rhs)?,
2017 |x| scalar * x,
2018 )
2019 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
2020 return Ok(tensor_from_array(out));
2021 }
2022 if rhs_is_scalar && lhs_is_full_output {
2023 let scalar = typed_view_from_view("broadcast_multiply", rhs)?.get(&[]);
2024 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs_shape) }?;
2026 map_into(
2027 &mut out.view_mut(),
2028 &typed_view_from_view("broadcast_multiply", lhs)?,
2029 |x| x * scalar,
2030 )
2031 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
2032 return Ok(tensor_from_array(out));
2033 }
2034
2035 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs_shape) }?;
2037 let lhs_view = typed_view_from_view("broadcast_multiply", lhs)?;
2038 let rhs_view = typed_view_from_view("broadcast_multiply", rhs)?;
2039 broadcast_mul_into(
2040 &mut out.view_mut(),
2041 &lhs_view,
2042 lhs_dims,
2043 &rhs_view,
2044 rhs_dims,
2045 )
2046 .map_err(|err| crate::Error::backend_failure("broadcast_multiply", err))?;
2047 Ok(tensor_from_array(out))
2048}
2049
2050fn filled_broadcast_multiply_tensor<T>(
2051 buffers: &mut BufferPool,
2052 shape: &[usize],
2053 fill: T,
2054) -> crate::Result<TypedTensor<T>>
2055where
2056 T: Copy + Clone + PoolScalar + 'static,
2057{
2058 let len = shape.iter().try_fold(1usize, |acc, &dim| {
2059 acc.checked_mul(dim).ok_or_else(|| {
2060 crate::Error::backend_failure("broadcast_multiply", "output shape size overflows usize")
2061 })
2062 })?;
2063 let mut data = unsafe { T::pool_acquire(buffers, len) };
2065 data.fill(fill);
2066 TypedTensor::from_vec_col_major(shape.to_vec(), data)
2067}
2068
2069#[allow(clippy::too_many_arguments)]
2070pub(crate) fn broadcast_multiply_read_with_pool(
2071 buffers: &mut BufferPool,
2072 lhs: TensorRead<'_>,
2073 lhs_shape: &[usize],
2074 lhs_dims: &[usize],
2075 rhs: TensorRead<'_>,
2076 rhs_shape: &[usize],
2077 rhs_dims: &[usize],
2078) -> crate::Result<Option<Tensor>> {
2079 let lhs = read_as_cpu_view(lhs);
2080 let rhs = read_as_cpu_view(rhs);
2081
2082 macro_rules! dispatch {
2083 ($variant:ident, $lhs:expr, $rhs:expr) => {{
2084 if let Some(out) = try_outer_product_with_pool(
2085 buffers, &$lhs, lhs_shape, lhs_dims, &$rhs, rhs_shape, rhs_dims,
2086 )? {
2087 return Ok(Some(Tensor::$variant(out)));
2088 }
2089 Ok(Some(Tensor::$variant(typed_broadcast_mul_view_with_pool(
2090 buffers, &$lhs, lhs_shape, lhs_dims, &$rhs, rhs_shape, rhs_dims,
2091 )?)))
2092 }};
2093 }
2094
2095 match (lhs, rhs) {
2096 (CpuReadView::F32(lhs), CpuReadView::F32(rhs)) => dispatch!(F32, lhs, rhs),
2097 (CpuReadView::F64(lhs), CpuReadView::F64(rhs)) => dispatch!(F64, lhs, rhs),
2098 (CpuReadView::I32(lhs), CpuReadView::I32(rhs)) => dispatch!(I32, lhs, rhs),
2099 (CpuReadView::I64(lhs), CpuReadView::I64(rhs)) => dispatch!(I64, lhs, rhs),
2100 (CpuReadView::C32(lhs), CpuReadView::C32(rhs)) => dispatch!(C32, lhs, rhs),
2101 (CpuReadView::C64(lhs), CpuReadView::C64(rhs)) => dispatch!(C64, lhs, rhs),
2102 _ => Ok(None),
2103 }
2104}
2105
2106#[allow(clippy::too_many_arguments)]
2107pub(crate) fn broadcast_multiply_value_with_pool(
2108 buffers: &mut BufferPool,
2109 lhs: TensorRead<'_>,
2110 lhs_shape: &[usize],
2111 lhs_dims: &[usize],
2112 rhs: TensorRead<'_>,
2113 rhs_shape: &[usize],
2114 rhs_dims: &[usize],
2115) -> crate::Result<Option<TensorValue>> {
2116 let lhs_view = read_as_cpu_view(lhs.clone());
2117 let rhs_view = read_as_cpu_view(rhs.clone());
2118
2119 macro_rules! dispatch_lazy {
2120 ($variant:ident, $lhs:expr, $rhs:expr) => {{
2121 if let Some(out) = try_lazy_outer_product_with_pool(
2122 buffers, &$lhs, lhs_shape, lhs_dims, &$rhs, rhs_shape, rhs_dims,
2123 )? {
2124 return Ok(Some(lazy_outer_product_value(
2125 Tensor::$variant(out.base),
2126 out.shape,
2127 out.strides,
2128 )?));
2129 }
2130 }};
2131 }
2132
2133 match (lhs_view, rhs_view) {
2134 (CpuReadView::F32(lhs_view), CpuReadView::F32(rhs_view)) => {
2135 dispatch_lazy!(F32, lhs_view, rhs_view);
2136 }
2137 (CpuReadView::F64(lhs_view), CpuReadView::F64(rhs_view)) => {
2138 dispatch_lazy!(F64, lhs_view, rhs_view);
2139 }
2140 (CpuReadView::I32(lhs_view), CpuReadView::I32(rhs_view)) => {
2141 dispatch_lazy!(I32, lhs_view, rhs_view);
2142 }
2143 (CpuReadView::I64(lhs_view), CpuReadView::I64(rhs_view)) => {
2144 dispatch_lazy!(I64, lhs_view, rhs_view);
2145 }
2146 (CpuReadView::C32(lhs_view), CpuReadView::C32(rhs_view)) => {
2147 dispatch_lazy!(C32, lhs_view, rhs_view);
2148 }
2149 (CpuReadView::C64(lhs_view), CpuReadView::C64(rhs_view)) => {
2150 dispatch_lazy!(C64, lhs_view, rhs_view);
2151 }
2152 _ => {}
2153 }
2154
2155 broadcast_multiply_read_with_pool(buffers, lhs, lhs_shape, lhs_dims, rhs, rhs_shape, rhs_dims)
2156 .map(|tensor| tensor.map(TensorValue::from_tensor))
2157}
2158
2159pub fn div(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2174 with_local_pool(|buffers| div_with_pool(buffers, lhs, rhs))
2175}
2176
2177pub(crate) fn div_with_pool(
2178 buffers: &mut BufferPool,
2179 lhs: &Tensor,
2180 rhs: &Tensor,
2181) -> crate::Result<Tensor> {
2182 match (lhs, rhs) {
2183 (Tensor::F32(a), Tensor::F32(b)) => Ok(Tensor::F32(typed_div_with_pool(buffers, a, b)?)),
2184 (Tensor::F64(a), Tensor::F64(b)) => Ok(Tensor::F64(typed_div_with_pool(buffers, a, b)?)),
2185 (Tensor::I32(a), Tensor::I32(b)) => {
2186 Ok(Tensor::I32(typed_integer_div_with_pool(buffers, a, b)?))
2187 }
2188 (Tensor::I64(a), Tensor::I64(b)) => {
2189 Ok(Tensor::I64(typed_integer_div_with_pool(buffers, a, b)?))
2190 }
2191 (Tensor::C32(a), Tensor::C32(b)) => Ok(Tensor::C32(typed_div_with_pool(buffers, a, b)?)),
2192 (Tensor::C64(a), Tensor::C64(b)) => Ok(Tensor::C64(typed_div_with_pool(buffers, a, b)?)),
2193 (Tensor::F32(a), Tensor::C32(b)) if a.shape().is_empty() => {
2194 let scalar = complex_scalar_tensor(typed_host_data("div", a)?[0])?;
2195 Ok(Tensor::C32(typed_div_with_pool(buffers, &scalar, b)?))
2196 }
2197 (Tensor::C32(a), Tensor::F32(b)) if b.shape().is_empty() => {
2198 let scalar = complex_scalar_tensor(typed_host_data("div", b)?[0])?;
2199 Ok(Tensor::C32(typed_div_with_pool(buffers, a, &scalar)?))
2200 }
2201 (Tensor::F64(a), Tensor::C64(b)) if a.shape().is_empty() => {
2202 let scalar = complex_scalar_tensor(typed_host_data("div", a)?[0])?;
2203 Ok(Tensor::C64(typed_div_with_pool(buffers, &scalar, b)?))
2204 }
2205 (Tensor::C64(a), Tensor::F64(b)) if b.shape().is_empty() => {
2206 let scalar = complex_scalar_tensor(typed_host_data("div", b)?[0])?;
2207 Ok(Tensor::C64(typed_div_with_pool(buffers, a, &scalar)?))
2208 }
2209 _ => Err(crate::Error::DTypeMismatch {
2210 op: "div",
2211 lhs: lhs.dtype(),
2212 rhs: rhs.dtype(),
2213 }),
2214 }
2215}
2216
2217pub(crate) fn div_read_with_pool(
2218 buffers: &mut BufferPool,
2219 lhs: TensorRead<'_>,
2220 rhs: TensorRead<'_>,
2221) -> crate::Result<Tensor> {
2222 let lhs_dtype = lhs.dtype();
2223 let rhs_dtype = rhs.dtype();
2224 match (read_as_cpu_view(lhs), read_as_cpu_view(rhs)) {
2225 (CpuReadView::F32(a), CpuReadView::F32(b)) => Ok(Tensor::F32(typed_binary_view_with_pool(
2226 "div",
2227 buffers,
2228 &a,
2229 &b,
2230 |x, y| x / y,
2231 )?)),
2232 (CpuReadView::F64(a), CpuReadView::F64(b)) => Ok(Tensor::F64(typed_binary_view_with_pool(
2233 "div",
2234 buffers,
2235 &a,
2236 &b,
2237 |x, y| x / y,
2238 )?)),
2239 (CpuReadView::I32(a), CpuReadView::I32(b)) => Ok(Tensor::I32(
2240 typed_integer_div_view_with_pool(buffers, &a, &b)?,
2241 )),
2242 (CpuReadView::I64(a), CpuReadView::I64(b)) => Ok(Tensor::I64(
2243 typed_integer_div_view_with_pool(buffers, &a, &b)?,
2244 )),
2245 (CpuReadView::C32(a), CpuReadView::C32(b)) => Ok(Tensor::C32(typed_binary_view_with_pool(
2246 "div",
2247 buffers,
2248 &a,
2249 &b,
2250 |x, y| x / y,
2251 )?)),
2252 (CpuReadView::C64(a), CpuReadView::C64(b)) => Ok(Tensor::C64(typed_binary_view_with_pool(
2253 "div",
2254 buffers,
2255 &a,
2256 &b,
2257 |x, y| x / y,
2258 )?)),
2259 (CpuReadView::F32(real), CpuReadView::C32(complex)) if real.shape().is_empty() => {
2260 let scalar = complex_scalar_tensor_from_view(&real)?;
2261 let scalar = scalar.as_view();
2262 Ok(Tensor::C32(typed_binary_view_with_pool(
2263 "div",
2264 buffers,
2265 &scalar,
2266 &complex,
2267 |x, y| x / y,
2268 )?))
2269 }
2270 (CpuReadView::C32(complex), CpuReadView::F32(real)) if real.shape().is_empty() => {
2271 let scalar = complex_scalar_tensor_from_view(&real)?;
2272 let scalar = scalar.as_view();
2273 Ok(Tensor::C32(typed_binary_view_with_pool(
2274 "div",
2275 buffers,
2276 &complex,
2277 &scalar,
2278 |x, y| x / y,
2279 )?))
2280 }
2281 (CpuReadView::F64(real), CpuReadView::C64(complex)) if real.shape().is_empty() => {
2282 let scalar = complex_scalar_tensor_from_view(&real)?;
2283 let scalar = scalar.as_view();
2284 Ok(Tensor::C64(typed_binary_view_with_pool(
2285 "div",
2286 buffers,
2287 &scalar,
2288 &complex,
2289 |x, y| x / y,
2290 )?))
2291 }
2292 (CpuReadView::C64(complex), CpuReadView::F64(real)) if real.shape().is_empty() => {
2293 let scalar = complex_scalar_tensor_from_view(&real)?;
2294 let scalar = scalar.as_view();
2295 Ok(Tensor::C64(typed_binary_view_with_pool(
2296 "div",
2297 buffers,
2298 &complex,
2299 &scalar,
2300 |x, y| x / y,
2301 )?))
2302 }
2303 _ => Err(crate::Error::DTypeMismatch {
2304 op: "div",
2305 lhs: lhs_dtype,
2306 rhs: rhs_dtype,
2307 }),
2308 }
2309}
2310
2311pub fn rem(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2329 with_local_pool(|buffers| rem_with_pool(buffers, lhs, rhs))
2330}
2331
2332pub(crate) fn rem_with_pool(
2333 buffers: &mut BufferPool,
2334 lhs: &Tensor,
2335 rhs: &Tensor,
2336) -> crate::Result<Tensor> {
2337 match (lhs, rhs) {
2338 (Tensor::F32(a), Tensor::F32(b)) => Ok(Tensor::F32(typed_rem_with_pool(buffers, a, b)?)),
2339 (Tensor::F64(a), Tensor::F64(b)) => Ok(Tensor::F64(typed_rem_with_pool(buffers, a, b)?)),
2340 (Tensor::I32(a), Tensor::I32(b)) => {
2341 Ok(Tensor::I32(typed_integer_rem_with_pool(buffers, a, b)?))
2342 }
2343 (Tensor::I64(a), Tensor::I64(b)) => {
2344 Ok(Tensor::I64(typed_integer_rem_with_pool(buffers, a, b)?))
2345 }
2346 _ => Err(tensor_pair_error("rem", lhs, rhs)),
2347 }
2348}
2349
2350pub(crate) fn rem_read_with_pool(
2351 buffers: &mut BufferPool,
2352 lhs: TensorRead<'_>,
2353 rhs: TensorRead<'_>,
2354) -> crate::Result<Tensor> {
2355 let lhs_dtype = lhs.dtype();
2356 let rhs_dtype = rhs.dtype();
2357 match (read_as_cpu_view(lhs), read_as_cpu_view(rhs)) {
2358 (CpuReadView::F32(a), CpuReadView::F32(b)) => Ok(Tensor::F32(typed_binary_view_with_pool(
2359 "rem",
2360 buffers,
2361 &a,
2362 &b,
2363 |x, y| x % y,
2364 )?)),
2365 (CpuReadView::F64(a), CpuReadView::F64(b)) => Ok(Tensor::F64(typed_binary_view_with_pool(
2366 "rem",
2367 buffers,
2368 &a,
2369 &b,
2370 |x, y| x % y,
2371 )?)),
2372 (CpuReadView::I32(a), CpuReadView::I32(b)) => Ok(Tensor::I32(
2373 typed_integer_rem_view_with_pool(buffers, &a, &b)?,
2374 )),
2375 (CpuReadView::I64(a), CpuReadView::I64(b)) => Ok(Tensor::I64(
2376 typed_integer_rem_view_with_pool(buffers, &a, &b)?,
2377 )),
2378 _ => Err(dtype_pair_error("rem", lhs_dtype, rhs_dtype)),
2379 }
2380}
2381
2382pub fn neg(input: &Tensor) -> crate::Result<Tensor> {
2396 with_local_pool(|buffers| neg_with_pool(buffers, input))
2397}
2398
2399pub(crate) fn neg_with_pool(buffers: &mut BufferPool, input: &Tensor) -> crate::Result<Tensor> {
2400 match input {
2401 Tensor::F32(t) => Ok(Tensor::F32(typed_neg_with_pool(buffers, t)?)),
2402 Tensor::F64(t) => Ok(Tensor::F64(typed_neg_with_pool(buffers, t)?)),
2403 Tensor::I32(t) => Ok(Tensor::I32(typed_wrapping_neg_with_pool(buffers, t)?)),
2404 Tensor::I64(t) => Ok(Tensor::I64(typed_wrapping_neg_with_pool(buffers, t)?)),
2405 Tensor::Bool(_) => Err(crate::Error::backend_failure(
2406 "neg",
2407 format!("unsupported dtype {:?}", input.dtype()),
2408 )),
2409 Tensor::C32(t) => Ok(Tensor::C32(typed_neg_with_pool(buffers, t)?)),
2410 Tensor::C64(t) => Ok(Tensor::C64(typed_neg_with_pool(buffers, t)?)),
2411 }
2412}
2413
2414pub(crate) fn neg_read_with_pool(
2415 buffers: &mut BufferPool,
2416 input: TensorRead<'_>,
2417) -> crate::Result<Tensor> {
2418 let dtype = input.dtype();
2419 match read_as_cpu_view(input) {
2420 CpuReadView::F32(t) => Ok(Tensor::F32(typed_unary_view_with_pool(
2421 "neg",
2422 buffers,
2423 &t,
2424 |x| -x,
2425 )?)),
2426 CpuReadView::F64(t) => Ok(Tensor::F64(typed_unary_view_with_pool(
2427 "neg",
2428 buffers,
2429 &t,
2430 |x| -x,
2431 )?)),
2432 CpuReadView::I32(t) => Ok(Tensor::I32(typed_unary_view_with_pool(
2433 "neg",
2434 buffers,
2435 &t,
2436 |x| x.wrapping_neg_elem(),
2437 )?)),
2438 CpuReadView::I64(t) => Ok(Tensor::I64(typed_unary_view_with_pool(
2439 "neg",
2440 buffers,
2441 &t,
2442 |x| x.wrapping_neg_elem(),
2443 )?)),
2444 CpuReadView::C32(t) => Ok(Tensor::C32(typed_unary_view_with_pool(
2445 "neg",
2446 buffers,
2447 &t,
2448 |x| -x,
2449 )?)),
2450 CpuReadView::C64(t) => Ok(Tensor::C64(typed_unary_view_with_pool(
2451 "neg",
2452 buffers,
2453 &t,
2454 |x| -x,
2455 )?)),
2456 _ => Err(crate::Error::backend_failure(
2457 "neg",
2458 format!("unsupported dtype {dtype:?}"),
2459 )),
2460 }
2461}
2462
2463pub fn conj(input: &Tensor) -> crate::Result<Tensor> {
2478 with_local_pool(|buffers| conj_with_pool(buffers, input))
2479}
2480
2481pub(crate) fn conj_with_pool(buffers: &mut BufferPool, input: &Tensor) -> crate::Result<Tensor> {
2482 match input {
2483 Tensor::F32(t) => Ok(Tensor::F32(typed_conj_with_pool(buffers, t)?)),
2484 Tensor::F64(t) => Ok(Tensor::F64(typed_conj_with_pool(buffers, t)?)),
2485 Tensor::I32(_) | Tensor::I64(_) | Tensor::Bool(_) => Err(crate::Error::backend_failure(
2486 "conj",
2487 format!("unsupported dtype {:?}", input.dtype()),
2488 )),
2489 Tensor::C32(t) => Ok(Tensor::C32(typed_conj_with_pool(buffers, t)?)),
2490 Tensor::C64(t) => Ok(Tensor::C64(typed_conj_with_pool(buffers, t)?)),
2491 }
2492}
2493
2494pub(crate) fn conj_read_with_pool(
2495 buffers: &mut BufferPool,
2496 input: TensorRead<'_>,
2497) -> crate::Result<Tensor> {
2498 let dtype = input.dtype();
2499 match read_as_cpu_view(input) {
2500 CpuReadView::F32(t) => Ok(Tensor::F32(typed_unary_view_with_pool(
2501 "conj",
2502 buffers,
2503 &t,
2504 |x| x.conj_elem(),
2505 )?)),
2506 CpuReadView::F64(t) => Ok(Tensor::F64(typed_unary_view_with_pool(
2507 "conj",
2508 buffers,
2509 &t,
2510 |x| x.conj_elem(),
2511 )?)),
2512 CpuReadView::C32(t) => Ok(Tensor::C32(typed_unary_view_with_pool(
2513 "conj",
2514 buffers,
2515 &t,
2516 |x| x.conj_elem(),
2517 )?)),
2518 CpuReadView::C64(t) => Ok(Tensor::C64(typed_unary_view_with_pool(
2519 "conj",
2520 buffers,
2521 &t,
2522 |x| x.conj_elem(),
2523 )?)),
2524 _ => Err(crate::Error::backend_failure(
2525 "conj",
2526 format!("unsupported dtype {dtype:?}"),
2527 )),
2528 }
2529}
2530
2531pub fn abs(input: &Tensor) -> crate::Result<Tensor> {
2547 with_local_pool(|buffers| abs_with_pool(buffers, input))
2548}
2549
2550pub(crate) fn abs_with_pool(buffers: &mut BufferPool, input: &Tensor) -> crate::Result<Tensor> {
2551 match input {
2552 Tensor::F32(t) => Ok(Tensor::F32(typed_abs_with_pool(buffers, t)?)),
2553 Tensor::F64(t) => Ok(Tensor::F64(typed_abs_with_pool(buffers, t)?)),
2554 Tensor::I32(t) => Ok(Tensor::I32(typed_wrapping_abs_with_pool(buffers, t)?)),
2555 Tensor::I64(t) => Ok(Tensor::I64(typed_wrapping_abs_with_pool(buffers, t)?)),
2556 Tensor::Bool(_) => Err(crate::Error::backend_failure(
2557 "abs",
2558 format!("unsupported dtype {:?}", input.dtype()),
2559 )),
2560 Tensor::C32(t) => Ok(Tensor::F32(typed_complex_abs_with_pool(buffers, t)?)),
2561 Tensor::C64(t) => Ok(Tensor::F64(typed_complex_abs_with_pool(buffers, t)?)),
2562 }
2563}
2564
2565pub(crate) fn abs_read_with_pool(
2566 buffers: &mut BufferPool,
2567 input: TensorRead<'_>,
2568) -> crate::Result<Tensor> {
2569 let dtype = input.dtype();
2570 match read_as_cpu_view(input) {
2571 CpuReadView::F32(t) => Ok(Tensor::F32(typed_unary_view_with_pool(
2572 "abs",
2573 buffers,
2574 &t,
2575 |x| x.abs_elem(),
2576 )?)),
2577 CpuReadView::F64(t) => Ok(Tensor::F64(typed_unary_view_with_pool(
2578 "abs",
2579 buffers,
2580 &t,
2581 |x| x.abs_elem(),
2582 )?)),
2583 CpuReadView::I32(t) => Ok(Tensor::I32(typed_unary_view_with_pool(
2584 "abs",
2585 buffers,
2586 &t,
2587 |x| x.wrapping_abs_elem(),
2588 )?)),
2589 CpuReadView::I64(t) => Ok(Tensor::I64(typed_unary_view_with_pool(
2590 "abs",
2591 buffers,
2592 &t,
2593 |x| x.wrapping_abs_elem(),
2594 )?)),
2595 CpuReadView::C32(t) => Ok(Tensor::F32(typed_complex_abs_view_with_pool(buffers, &t)?)),
2596 CpuReadView::C64(t) => Ok(Tensor::F64(typed_complex_abs_view_with_pool(buffers, &t)?)),
2597 _ => Err(crate::Error::backend_failure(
2598 "abs",
2599 format!("unsupported dtype {dtype:?}"),
2600 )),
2601 }
2602}
2603
2604pub fn sign(input: &Tensor) -> crate::Result<Tensor> {
2618 with_local_pool(|buffers| sign_with_pool(buffers, input))
2619}
2620
2621pub(crate) fn sign_with_pool(buffers: &mut BufferPool, input: &Tensor) -> crate::Result<Tensor> {
2622 match input {
2623 Tensor::F32(t) => Ok(Tensor::F32(typed_sign_with_pool(buffers, t)?)),
2624 Tensor::F64(t) => Ok(Tensor::F64(typed_sign_with_pool(buffers, t)?)),
2625 Tensor::I32(t) => Ok(Tensor::I32(typed_integer_sign_with_pool(buffers, t)?)),
2626 Tensor::I64(t) => Ok(Tensor::I64(typed_integer_sign_with_pool(buffers, t)?)),
2627 Tensor::Bool(_) => Err(crate::Error::backend_failure(
2628 "sign",
2629 format!("unsupported dtype {:?}", input.dtype()),
2630 )),
2631 Tensor::C32(t) => Ok(Tensor::C32(typed_sign_with_pool(buffers, t)?)),
2632 Tensor::C64(t) => Ok(Tensor::C64(typed_sign_with_pool(buffers, t)?)),
2633 }
2634}
2635
2636pub(crate) fn sign_read_with_pool(
2637 buffers: &mut BufferPool,
2638 input: TensorRead<'_>,
2639) -> crate::Result<Tensor> {
2640 let dtype = input.dtype();
2641 match read_as_cpu_view(input) {
2642 CpuReadView::F32(t) => Ok(Tensor::F32(typed_unary_view_with_pool(
2643 "sign",
2644 buffers,
2645 &t,
2646 |x| x.sign_elem(),
2647 )?)),
2648 CpuReadView::F64(t) => Ok(Tensor::F64(typed_unary_view_with_pool(
2649 "sign",
2650 buffers,
2651 &t,
2652 |x| x.sign_elem(),
2653 )?)),
2654 CpuReadView::I32(t) => Ok(Tensor::I32(typed_unary_view_with_pool(
2655 "sign",
2656 buffers,
2657 &t,
2658 |x| x.signum_elem(),
2659 )?)),
2660 CpuReadView::I64(t) => Ok(Tensor::I64(typed_unary_view_with_pool(
2661 "sign",
2662 buffers,
2663 &t,
2664 |x| x.signum_elem(),
2665 )?)),
2666 CpuReadView::C32(t) => Ok(Tensor::C32(typed_unary_view_with_pool(
2667 "sign",
2668 buffers,
2669 &t,
2670 |x| x.sign_elem(),
2671 )?)),
2672 CpuReadView::C64(t) => Ok(Tensor::C64(typed_unary_view_with_pool(
2673 "sign",
2674 buffers,
2675 &t,
2676 |x| x.sign_elem(),
2677 )?)),
2678 _ => Err(crate::Error::backend_failure(
2679 "sign",
2680 format!("unsupported dtype {dtype:?}"),
2681 )),
2682 }
2683}
2684
2685pub fn maximum(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2700 with_local_pool(|buffers| maximum_with_pool(buffers, lhs, rhs))
2701}
2702
2703pub(crate) fn maximum_with_pool(
2704 buffers: &mut BufferPool,
2705 lhs: &Tensor,
2706 rhs: &Tensor,
2707) -> crate::Result<Tensor> {
2708 reject_complex_ordered_dtypes("maximum", &[lhs.dtype(), rhs.dtype()])?;
2709
2710 match (lhs, rhs) {
2711 (Tensor::F32(a), Tensor::F32(b)) => {
2712 Ok(Tensor::F32(typed_maximum_with_pool(buffers, a, b)?))
2713 }
2714 (Tensor::F64(a), Tensor::F64(b)) => {
2715 Ok(Tensor::F64(typed_maximum_with_pool(buffers, a, b)?))
2716 }
2717 (Tensor::I32(a), Tensor::I32(b)) => {
2718 Ok(Tensor::I32(typed_maximum_with_pool(buffers, a, b)?))
2719 }
2720 (Tensor::I64(a), Tensor::I64(b)) => {
2721 Ok(Tensor::I64(typed_maximum_with_pool(buffers, a, b)?))
2722 }
2723 _ => Err(tensor_pair_error("maximum", lhs, rhs)),
2724 }
2725}
2726
2727pub(crate) fn maximum_read_with_pool(
2728 buffers: &mut BufferPool,
2729 lhs: TensorRead<'_>,
2730 rhs: TensorRead<'_>,
2731) -> crate::Result<Tensor> {
2732 let lhs_dtype = lhs.dtype();
2733 let rhs_dtype = rhs.dtype();
2734 reject_complex_ordered_dtypes("maximum", &[lhs_dtype, rhs_dtype])?;
2735
2736 match (read_as_cpu_view(lhs), read_as_cpu_view(rhs)) {
2737 (CpuReadView::F32(a), CpuReadView::F32(b)) => Ok(Tensor::F32(
2738 typed_same_shape_binary_view_with_pool("maximum", buffers, &a, &b, |x, y| {
2739 x.max_elem(y)
2740 })?,
2741 )),
2742 (CpuReadView::F64(a), CpuReadView::F64(b)) => Ok(Tensor::F64(
2743 typed_same_shape_binary_view_with_pool("maximum", buffers, &a, &b, |x, y| {
2744 x.max_elem(y)
2745 })?,
2746 )),
2747 (CpuReadView::I32(a), CpuReadView::I32(b)) => Ok(Tensor::I32(
2748 typed_same_shape_binary_view_with_pool("maximum", buffers, &a, &b, |x, y| {
2749 x.max_elem(y)
2750 })?,
2751 )),
2752 (CpuReadView::I64(a), CpuReadView::I64(b)) => Ok(Tensor::I64(
2753 typed_same_shape_binary_view_with_pool("maximum", buffers, &a, &b, |x, y| {
2754 x.max_elem(y)
2755 })?,
2756 )),
2757 _ => Err(dtype_pair_error("maximum", lhs_dtype, rhs_dtype)),
2758 }
2759}
2760
2761pub fn minimum(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
2776 with_local_pool(|buffers| minimum_with_pool(buffers, lhs, rhs))
2777}
2778
2779pub(crate) fn minimum_with_pool(
2780 buffers: &mut BufferPool,
2781 lhs: &Tensor,
2782 rhs: &Tensor,
2783) -> crate::Result<Tensor> {
2784 reject_complex_ordered_dtypes("minimum", &[lhs.dtype(), rhs.dtype()])?;
2785
2786 match (lhs, rhs) {
2787 (Tensor::F32(a), Tensor::F32(b)) => {
2788 Ok(Tensor::F32(typed_minimum_with_pool(buffers, a, b)?))
2789 }
2790 (Tensor::F64(a), Tensor::F64(b)) => {
2791 Ok(Tensor::F64(typed_minimum_with_pool(buffers, a, b)?))
2792 }
2793 (Tensor::I32(a), Tensor::I32(b)) => {
2794 Ok(Tensor::I32(typed_minimum_with_pool(buffers, a, b)?))
2795 }
2796 (Tensor::I64(a), Tensor::I64(b)) => {
2797 Ok(Tensor::I64(typed_minimum_with_pool(buffers, a, b)?))
2798 }
2799 _ => Err(tensor_pair_error("minimum", lhs, rhs)),
2800 }
2801}
2802
2803pub(crate) fn minimum_read_with_pool(
2804 buffers: &mut BufferPool,
2805 lhs: TensorRead<'_>,
2806 rhs: TensorRead<'_>,
2807) -> crate::Result<Tensor> {
2808 let lhs_dtype = lhs.dtype();
2809 let rhs_dtype = rhs.dtype();
2810 reject_complex_ordered_dtypes("minimum", &[lhs_dtype, rhs_dtype])?;
2811
2812 match (read_as_cpu_view(lhs), read_as_cpu_view(rhs)) {
2813 (CpuReadView::F32(a), CpuReadView::F32(b)) => Ok(Tensor::F32(
2814 typed_same_shape_binary_view_with_pool("minimum", buffers, &a, &b, |x, y| {
2815 x.min_elem(y)
2816 })?,
2817 )),
2818 (CpuReadView::F64(a), CpuReadView::F64(b)) => Ok(Tensor::F64(
2819 typed_same_shape_binary_view_with_pool("minimum", buffers, &a, &b, |x, y| {
2820 x.min_elem(y)
2821 })?,
2822 )),
2823 (CpuReadView::I32(a), CpuReadView::I32(b)) => Ok(Tensor::I32(
2824 typed_same_shape_binary_view_with_pool("minimum", buffers, &a, &b, |x, y| {
2825 x.min_elem(y)
2826 })?,
2827 )),
2828 (CpuReadView::I64(a), CpuReadView::I64(b)) => Ok(Tensor::I64(
2829 typed_same_shape_binary_view_with_pool("minimum", buffers, &a, &b, |x, y| {
2830 x.min_elem(y)
2831 })?,
2832 )),
2833 _ => Err(dtype_pair_error("minimum", lhs_dtype, rhs_dtype)),
2834 }
2835}
2836
2837pub fn compare(lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> crate::Result<Tensor> {
2852 with_local_pool(|buffers| compare_with_pool(buffers, lhs, rhs, dir))
2853}
2854
2855pub(crate) fn compare_with_pool(
2856 buffers: &mut BufferPool,
2857 lhs: &Tensor,
2858 rhs: &Tensor,
2859 dir: &CompareDir,
2860) -> crate::Result<Tensor> {
2861 reject_complex_ordered_dtypes("compare", &[lhs.dtype(), rhs.dtype()])?;
2862
2863 match (lhs, rhs) {
2864 (Tensor::F32(a), Tensor::F32(b)) => {
2865 Ok(Tensor::Bool(typed_compare_with_pool(buffers, a, b, dir)?))
2866 }
2867 (Tensor::F64(a), Tensor::F64(b)) => {
2868 Ok(Tensor::Bool(typed_compare_with_pool(buffers, a, b, dir)?))
2869 }
2870 (Tensor::I32(a), Tensor::I32(b)) => {
2871 Ok(Tensor::Bool(typed_compare_with_pool(buffers, a, b, dir)?))
2872 }
2873 (Tensor::I64(a), Tensor::I64(b)) => {
2874 Ok(Tensor::Bool(typed_compare_with_pool(buffers, a, b, dir)?))
2875 }
2876 (Tensor::Bool(a), Tensor::Bool(b)) => {
2877 Ok(Tensor::Bool(typed_compare_with_pool(buffers, a, b, dir)?))
2878 }
2879 _ => Err(crate::Error::DTypeMismatch {
2880 op: "compare",
2881 lhs: lhs.dtype(),
2882 rhs: rhs.dtype(),
2883 }),
2884 }
2885}
2886
2887pub(crate) fn compare_read_with_pool(
2888 buffers: &mut BufferPool,
2889 lhs: TensorRead<'_>,
2890 rhs: TensorRead<'_>,
2891 dir: &CompareDir,
2892) -> crate::Result<Tensor> {
2893 let lhs_dtype = lhs.dtype();
2894 let rhs_dtype = rhs.dtype();
2895 reject_complex_ordered_dtypes("compare", &[lhs_dtype, rhs_dtype])?;
2896
2897 match (read_as_cpu_view(lhs), read_as_cpu_view(rhs)) {
2898 (CpuReadView::F32(a), CpuReadView::F32(b)) => Ok(Tensor::Bool(
2899 typed_same_shape_binary_view_with_pool("compare", buffers, &a, &b, |x, y| {
2900 x.compare_elem(y, dir)
2901 })?,
2902 )),
2903 (CpuReadView::F64(a), CpuReadView::F64(b)) => Ok(Tensor::Bool(
2904 typed_same_shape_binary_view_with_pool("compare", buffers, &a, &b, |x, y| {
2905 x.compare_elem(y, dir)
2906 })?,
2907 )),
2908 (CpuReadView::I32(a), CpuReadView::I32(b)) => Ok(Tensor::Bool(
2909 typed_same_shape_binary_view_with_pool("compare", buffers, &a, &b, |x, y| {
2910 x.compare_elem(y, dir)
2911 })?,
2912 )),
2913 (CpuReadView::I64(a), CpuReadView::I64(b)) => Ok(Tensor::Bool(
2914 typed_same_shape_binary_view_with_pool("compare", buffers, &a, &b, |x, y| {
2915 x.compare_elem(y, dir)
2916 })?,
2917 )),
2918 (CpuReadView::Bool(a), CpuReadView::Bool(b)) => Ok(Tensor::Bool(
2919 typed_same_shape_binary_view_with_pool("compare", buffers, &a, &b, |x, y| {
2920 x.compare_elem(y, dir)
2921 })?,
2922 )),
2923 _ => Err(crate::Error::DTypeMismatch {
2924 op: "compare",
2925 lhs: lhs_dtype,
2926 rhs: rhs_dtype,
2927 }),
2928 }
2929}
2930
2931pub fn select(pred: &Tensor, on_true: &Tensor, on_false: &Tensor) -> crate::Result<Tensor> {
2947 with_local_pool(|buffers| select_with_pool(buffers, pred, on_true, on_false))
2948}
2949
2950pub(crate) fn select_with_pool(
2951 buffers: &mut BufferPool,
2952 pred: &Tensor,
2953 on_true: &Tensor,
2954 on_false: &Tensor,
2955) -> crate::Result<Tensor> {
2956 match (pred, on_true, on_false) {
2957 (Tensor::Bool(p), Tensor::F32(t), Tensor::F32(f)) => {
2958 Ok(Tensor::F32(typed_select_with_pool(buffers, p, t, f)?))
2959 }
2960 (Tensor::Bool(p), Tensor::F64(t), Tensor::F64(f)) => {
2961 Ok(Tensor::F64(typed_select_with_pool(buffers, p, t, f)?))
2962 }
2963 (Tensor::Bool(p), Tensor::I32(t), Tensor::I32(f)) => {
2964 Ok(Tensor::I32(typed_select_with_pool(buffers, p, t, f)?))
2965 }
2966 (Tensor::Bool(p), Tensor::I64(t), Tensor::I64(f)) => {
2967 Ok(Tensor::I64(typed_select_with_pool(buffers, p, t, f)?))
2968 }
2969 (Tensor::Bool(p), Tensor::Bool(t), Tensor::Bool(f)) => {
2970 Ok(Tensor::Bool(typed_select_with_pool(buffers, p, t, f)?))
2971 }
2972 (Tensor::Bool(p), Tensor::C32(t), Tensor::C32(f)) => {
2973 Ok(Tensor::C32(typed_select_with_pool(buffers, p, t, f)?))
2974 }
2975 (Tensor::Bool(p), Tensor::C64(t), Tensor::C64(f)) => {
2976 Ok(Tensor::C64(typed_select_with_pool(buffers, p, t, f)?))
2977 }
2978 (Tensor::Bool(_), _, _) => Err(crate::Error::DTypeMismatch {
2979 op: "select",
2980 lhs: on_true.dtype(),
2981 rhs: on_false.dtype(),
2982 }),
2983 _ => Err(crate::Error::DTypeMismatch {
2984 op: "select",
2985 lhs: pred.dtype(),
2986 rhs: crate::DType::Bool,
2987 }),
2988 }
2989}
2990
2991pub(crate) fn select_read_with_pool(
2992 buffers: &mut BufferPool,
2993 pred: TensorRead<'_>,
2994 on_true: TensorRead<'_>,
2995 on_false: TensorRead<'_>,
2996) -> crate::Result<Tensor> {
2997 let pred_dtype = pred.dtype();
2998 let true_dtype = on_true.dtype();
2999 let false_dtype = on_false.dtype();
3000 match (
3001 read_as_cpu_view(pred),
3002 read_as_cpu_view(on_true),
3003 read_as_cpu_view(on_false),
3004 ) {
3005 (CpuReadView::Bool(p), CpuReadView::F32(t), CpuReadView::F32(f)) => Ok(Tensor::F32(
3006 typed_select_view_with_pool(buffers, &p, &t, &f)?,
3007 )),
3008 (CpuReadView::Bool(p), CpuReadView::F64(t), CpuReadView::F64(f)) => Ok(Tensor::F64(
3009 typed_select_view_with_pool(buffers, &p, &t, &f)?,
3010 )),
3011 (CpuReadView::Bool(p), CpuReadView::I32(t), CpuReadView::I32(f)) => Ok(Tensor::I32(
3012 typed_select_view_with_pool(buffers, &p, &t, &f)?,
3013 )),
3014 (CpuReadView::Bool(p), CpuReadView::I64(t), CpuReadView::I64(f)) => Ok(Tensor::I64(
3015 typed_select_view_with_pool(buffers, &p, &t, &f)?,
3016 )),
3017 (CpuReadView::Bool(p), CpuReadView::Bool(t), CpuReadView::Bool(f)) => Ok(Tensor::Bool(
3018 typed_select_view_with_pool(buffers, &p, &t, &f)?,
3019 )),
3020 (CpuReadView::Bool(p), CpuReadView::C32(t), CpuReadView::C32(f)) => Ok(Tensor::C32(
3021 typed_select_view_with_pool(buffers, &p, &t, &f)?,
3022 )),
3023 (CpuReadView::Bool(p), CpuReadView::C64(t), CpuReadView::C64(f)) => Ok(Tensor::C64(
3024 typed_select_view_with_pool(buffers, &p, &t, &f)?,
3025 )),
3026 (CpuReadView::Bool(_), _, _) => Err(crate::Error::DTypeMismatch {
3027 op: "select",
3028 lhs: true_dtype,
3029 rhs: false_dtype,
3030 }),
3031 _ => Err(crate::Error::DTypeMismatch {
3032 op: "select",
3033 lhs: pred_dtype,
3034 rhs: crate::DType::Bool,
3035 }),
3036 }
3037}
3038
3039pub fn clamp(input: &Tensor, lower: &Tensor, upper: &Tensor) -> crate::Result<Tensor> {
3055 with_local_pool(|buffers| clamp_with_pool(buffers, input, lower, upper))
3056}
3057
3058pub(crate) fn clamp_with_pool(
3059 buffers: &mut BufferPool,
3060 input: &Tensor,
3061 lower: &Tensor,
3062 upper: &Tensor,
3063) -> crate::Result<Tensor> {
3064 reject_complex_ordered_dtypes("clamp", &[input.dtype(), lower.dtype(), upper.dtype()])?;
3065
3066 dispatch_ternary_result_with_pool!("clamp", input, lower, upper, |x, lo, hi| {
3067 typed_clamp_with_pool(buffers, x, lo, hi)
3068 })
3069}
3070
3071pub(crate) fn clamp_read_with_pool(
3072 buffers: &mut BufferPool,
3073 input: TensorRead<'_>,
3074 lower: TensorRead<'_>,
3075 upper: TensorRead<'_>,
3076) -> crate::Result<Tensor> {
3077 let input_dtype = input.dtype();
3078 let lower_dtype = lower.dtype();
3079 let upper_dtype = upper.dtype();
3080 reject_complex_ordered_dtypes("clamp", &[input_dtype, lower_dtype, upper_dtype])?;
3081
3082 match (
3083 read_as_cpu_view(input),
3084 read_as_cpu_view(lower),
3085 read_as_cpu_view(upper),
3086 ) {
3087 (CpuReadView::F32(input), CpuReadView::F32(lower), CpuReadView::F32(upper)) => Ok(
3088 Tensor::F32(typed_clamp_view_with_pool(buffers, &input, &lower, &upper)?),
3089 ),
3090 (CpuReadView::F64(input), CpuReadView::F64(lower), CpuReadView::F64(upper)) => Ok(
3091 Tensor::F64(typed_clamp_view_with_pool(buffers, &input, &lower, &upper)?),
3092 ),
3093 _ => Err(crate::Error::DTypeMismatch {
3094 op: "clamp",
3095 lhs: input_dtype,
3096 rhs: lower_dtype,
3097 }),
3098 }
3099}
3100
3101pub(crate) fn typed_add_with_pool<T>(
3102 buffers: &mut BufferPool,
3103 lhs: &TypedTensor<T>,
3104 rhs: &TypedTensor<T>,
3105) -> crate::Result<TypedTensor<T>>
3106where
3107 T: Copy + Clone + Zero + Add<Output = T> + PoolScalar,
3108{
3109 if lhs.shape() == rhs.shape() {
3110 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3112 zip_map2_into(
3113 &mut out.view_mut(),
3114 &typed_view("add", lhs)?,
3115 &typed_view("add", rhs)?,
3116 |x, y| x + y,
3117 )
3118 .map_err(|err| crate::Error::backend_failure("add", err.to_string()))?;
3119 Ok(tensor_from_array(out))
3120 } else if lhs.shape().is_empty() {
3121 let scalar = typed_host_data("add", lhs)?[0];
3122 let mut out = unsafe { typed_array_uninit_from_pool(buffers, rhs.shape()) }?;
3124 map_into(&mut out.view_mut(), &typed_view("add", rhs)?, |x| {
3125 scalar + x
3126 })
3127 .map_err(|err| crate::Error::backend_failure("add", err.to_string()))?;
3128 Ok(tensor_from_array(out))
3129 } else if rhs.shape().is_empty() {
3130 let scalar = typed_host_data("add", rhs)?[0];
3131 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3133 map_into(&mut out.view_mut(), &typed_view("add", lhs)?, |x| {
3134 x + scalar
3135 })
3136 .map_err(|err| crate::Error::backend_failure("add", err.to_string()))?;
3137 Ok(tensor_from_array(out))
3138 } else {
3139 Err(crate::Error::ShapeMismatch {
3140 op: "add",
3141 lhs: lhs.shape().to_vec(),
3142 rhs: rhs.shape().to_vec(),
3143 })
3144 }
3145}
3146
3147fn typed_binary_with_pool<T>(
3148 op: &'static str,
3149 buffers: &mut BufferPool,
3150 lhs: &TypedTensor<T>,
3151 rhs: &TypedTensor<T>,
3152 f: impl Fn(T, T) -> T + Copy + Sync,
3153) -> crate::Result<TypedTensor<T>>
3154where
3155 T: Copy + PoolScalar + 'static,
3156{
3157 if lhs.shape() == rhs.shape() {
3158 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3160 zip_map2_into(
3161 &mut out.view_mut(),
3162 &typed_view(op, lhs)?,
3163 &typed_view(op, rhs)?,
3164 f,
3165 )
3166 .map_err(|err| crate::Error::backend_failure(op, err.to_string()))?;
3167 Ok(tensor_from_array(out))
3168 } else if lhs.shape().is_empty() {
3169 let scalar = typed_host_data(op, lhs)?[0];
3170 let mut out = unsafe { typed_array_uninit_from_pool(buffers, rhs.shape()) }?;
3172 map_into(&mut out.view_mut(), &typed_view(op, rhs)?, |x| f(scalar, x))
3173 .map_err(|err| crate::Error::backend_failure(op, err.to_string()))?;
3174 Ok(tensor_from_array(out))
3175 } else if rhs.shape().is_empty() {
3176 let scalar = typed_host_data(op, rhs)?[0];
3177 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3179 map_into(&mut out.view_mut(), &typed_view(op, lhs)?, |x| f(x, scalar))
3180 .map_err(|err| crate::Error::backend_failure(op, err.to_string()))?;
3181 Ok(tensor_from_array(out))
3182 } else {
3183 Err(crate::Error::ShapeMismatch {
3184 op,
3185 lhs: lhs.shape().to_vec(),
3186 rhs: rhs.shape().to_vec(),
3187 })
3188 }
3189}
3190
3191fn typed_wrapping_add_with_pool<T>(
3192 buffers: &mut BufferPool,
3193 lhs: &TypedTensor<T>,
3194 rhs: &TypedTensor<T>,
3195) -> crate::Result<TypedTensor<T>>
3196where
3197 T: WrappingIntegerElem,
3198{
3199 typed_binary_with_pool("add", buffers, lhs, rhs, |x, y| x.wrapping_add_elem(y))
3200}
3201
3202fn typed_wrapping_add_view_with_pool<T, L, R>(
3203 buffers: &mut BufferPool,
3204 lhs: &TypedTensorView<'_, T, L>,
3205 rhs: &TypedTensorView<'_, T, R>,
3206) -> crate::Result<TypedTensor<T>>
3207where
3208 T: WrappingIntegerElem,
3209 L: TensorRank,
3210 R: TensorRank,
3211{
3212 typed_binary_view_with_pool("add", buffers, lhs, rhs, |x, y| x.wrapping_add_elem(y))
3213}
3214
3215pub(crate) fn typed_add_view_with_pool<T, L, R>(
3216 buffers: &mut BufferPool,
3217 lhs: &TypedTensorView<'_, T, L>,
3218 rhs: &TypedTensorView<'_, T, R>,
3219) -> crate::Result<TypedTensor<T>>
3220where
3221 T: Copy + Clone + Zero + Add<Output = T> + PoolScalar + 'static,
3222 L: TensorRank,
3223 R: TensorRank,
3224{
3225 if lhs.shape() == rhs.shape() {
3226 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3228 zip_map2_into(
3229 &mut out.view_mut(),
3230 &typed_view_from_view("add", lhs)?,
3231 &typed_view_from_view("add", rhs)?,
3232 |x, y| x + y,
3233 )
3234 .map_err(|err| crate::Error::backend_failure("add", err.to_string()))?;
3235 Ok(tensor_from_array(out))
3236 } else if lhs.shape().is_empty() {
3237 let scalar = typed_view_from_view("add", lhs)?.get(&[]);
3238 let mut out = unsafe { typed_array_uninit_from_pool(buffers, rhs.shape()) }?;
3240 map_into(
3241 &mut out.view_mut(),
3242 &typed_view_from_view("add", rhs)?,
3243 |x| scalar + x,
3244 )
3245 .map_err(|err| crate::Error::backend_failure("add", err.to_string()))?;
3246 Ok(tensor_from_array(out))
3247 } else if rhs.shape().is_empty() {
3248 let scalar = typed_view_from_view("add", rhs)?.get(&[]);
3249 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3251 map_into(
3252 &mut out.view_mut(),
3253 &typed_view_from_view("add", lhs)?,
3254 |x| x + scalar,
3255 )
3256 .map_err(|err| crate::Error::backend_failure("add", err.to_string()))?;
3257 Ok(tensor_from_array(out))
3258 } else {
3259 Err(crate::Error::ShapeMismatch {
3260 op: "add",
3261 lhs: lhs.shape().to_vec(),
3262 rhs: rhs.shape().to_vec(),
3263 })
3264 }
3265}
3266
3267pub(crate) fn typed_sub_with_pool<T>(
3268 buffers: &mut BufferPool,
3269 lhs: &TypedTensor<T>,
3270 rhs: &TypedTensor<T>,
3271) -> crate::Result<TypedTensor<T>>
3272where
3273 T: Copy + PoolScalar + Sub<Output = T> + 'static,
3274{
3275 typed_binary_with_pool("sub", buffers, lhs, rhs, |x, y| x - y)
3276}
3277
3278fn typed_wrapping_sub_with_pool<T>(
3279 buffers: &mut BufferPool,
3280 lhs: &TypedTensor<T>,
3281 rhs: &TypedTensor<T>,
3282) -> crate::Result<TypedTensor<T>>
3283where
3284 T: WrappingIntegerElem,
3285{
3286 typed_binary_with_pool("sub", buffers, lhs, rhs, |x, y| x.wrapping_sub_elem(y))
3287}
3288
3289fn typed_wrapping_sub_view_with_pool<T, L, R>(
3290 buffers: &mut BufferPool,
3291 lhs: &TypedTensorView<'_, T, L>,
3292 rhs: &TypedTensorView<'_, T, R>,
3293) -> crate::Result<TypedTensor<T>>
3294where
3295 T: WrappingIntegerElem,
3296 L: TensorRank,
3297 R: TensorRank,
3298{
3299 typed_binary_view_with_pool("sub", buffers, lhs, rhs, |x, y| x.wrapping_sub_elem(y))
3300}
3301
3302pub(crate) fn typed_sub_view_with_pool<T, L, R>(
3303 buffers: &mut BufferPool,
3304 lhs: &TypedTensorView<'_, T, L>,
3305 rhs: &TypedTensorView<'_, T, R>,
3306) -> crate::Result<TypedTensor<T>>
3307where
3308 T: Copy + PoolScalar + Sub<Output = T> + 'static,
3309 L: TensorRank,
3310 R: TensorRank,
3311{
3312 typed_binary_view_with_pool("sub", buffers, lhs, rhs, |x, y| x - y)
3313}
3314
3315pub(crate) fn typed_mul_with_pool<T>(
3316 buffers: &mut BufferPool,
3317 lhs: &TypedTensor<T>,
3318 rhs: &TypedTensor<T>,
3319) -> crate::Result<TypedTensor<T>>
3320where
3321 T: Copy + Clone + Zero + Mul<Output = T> + PoolScalar + 'static,
3322{
3323 if lhs.shape() == rhs.shape() {
3324 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3326 mul_into(
3327 &mut out.view_mut(),
3328 &typed_view("mul", lhs)?,
3329 &typed_view("mul", rhs)?,
3330 )
3331 .map_err(|err| crate::Error::backend_failure("mul", err))?;
3332 Ok(tensor_from_array(out))
3333 } else if lhs.shape().is_empty() {
3334 let scalar = typed_host_data("mul", lhs)?[0];
3335 let mut out = unsafe { typed_array_uninit_from_pool(buffers, rhs.shape()) }?;
3337 map_into(&mut out.view_mut(), &typed_view("mul", rhs)?, |x| {
3338 scalar * x
3339 })
3340 .map_err(|err| crate::Error::backend_failure("mul", err))?;
3341 Ok(tensor_from_array(out))
3342 } else if rhs.shape().is_empty() {
3343 let scalar = typed_host_data("mul", rhs)?[0];
3344 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3346 map_into(&mut out.view_mut(), &typed_view("mul", lhs)?, |x| {
3347 x * scalar
3348 })
3349 .map_err(|err| crate::Error::backend_failure("mul", err))?;
3350 Ok(tensor_from_array(out))
3351 } else {
3352 Err(crate::Error::ShapeMismatch {
3353 op: "mul",
3354 lhs: lhs.shape().to_vec(),
3355 rhs: rhs.shape().to_vec(),
3356 })
3357 }
3358}
3359
3360fn typed_wrapping_mul_with_pool<T>(
3361 buffers: &mut BufferPool,
3362 lhs: &TypedTensor<T>,
3363 rhs: &TypedTensor<T>,
3364) -> crate::Result<TypedTensor<T>>
3365where
3366 T: WrappingIntegerElem,
3367{
3368 typed_binary_with_pool("mul", buffers, lhs, rhs, |x, y| x.wrapping_mul_elem(y))
3369}
3370
3371fn typed_wrapping_mul_view_with_pool<T, L, R>(
3372 buffers: &mut BufferPool,
3373 lhs: &TypedTensorView<'_, T, L>,
3374 rhs: &TypedTensorView<'_, T, R>,
3375) -> crate::Result<TypedTensor<T>>
3376where
3377 T: WrappingIntegerElem,
3378 L: TensorRank,
3379 R: TensorRank,
3380{
3381 typed_binary_view_with_pool("mul", buffers, lhs, rhs, |x, y| x.wrapping_mul_elem(y))
3382}
3383
3384pub(crate) fn typed_mul_view_with_pool<T, L, R>(
3385 buffers: &mut BufferPool,
3386 lhs: &TypedTensorView<'_, T, L>,
3387 rhs: &TypedTensorView<'_, T, R>,
3388) -> crate::Result<TypedTensor<T>>
3389where
3390 T: Copy + Clone + Zero + Mul<Output = T> + PoolScalar + 'static,
3391 L: TensorRank,
3392 R: TensorRank,
3393{
3394 if lhs.shape() == rhs.shape() {
3395 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3397 mul_into(
3398 &mut out.view_mut(),
3399 &typed_view_from_view("mul", lhs)?,
3400 &typed_view_from_view("mul", rhs)?,
3401 )
3402 .map_err(|err| crate::Error::backend_failure("mul", err))?;
3403 Ok(tensor_from_array(out))
3404 } else if lhs.shape().is_empty() {
3405 let scalar = typed_view_from_view("mul", lhs)?.get(&[]);
3406 let mut out = unsafe { typed_array_uninit_from_pool(buffers, rhs.shape()) }?;
3408 map_into(
3409 &mut out.view_mut(),
3410 &typed_view_from_view("mul", rhs)?,
3411 |x| scalar * x,
3412 )
3413 .map_err(|err| crate::Error::backend_failure("mul", err))?;
3414 Ok(tensor_from_array(out))
3415 } else if rhs.shape().is_empty() {
3416 let scalar = typed_view_from_view("mul", rhs)?.get(&[]);
3417 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3419 map_into(
3420 &mut out.view_mut(),
3421 &typed_view_from_view("mul", lhs)?,
3422 |x| x * scalar,
3423 )
3424 .map_err(|err| crate::Error::backend_failure("mul", err))?;
3425 Ok(tensor_from_array(out))
3426 } else {
3427 Err(crate::Error::ShapeMismatch {
3428 op: "mul",
3429 lhs: lhs.shape().to_vec(),
3430 rhs: rhs.shape().to_vec(),
3431 })
3432 }
3433}
3434
3435pub(crate) fn typed_div_with_pool<T>(
3436 buffers: &mut BufferPool,
3437 lhs: &TypedTensor<T>,
3438 rhs: &TypedTensor<T>,
3439) -> crate::Result<TypedTensor<T>>
3440where
3441 T: Copy + Clone + Zero + Div<Output = T> + PoolScalar,
3442{
3443 if lhs.shape() == rhs.shape() {
3444 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3446 zip_map2_into(
3447 &mut out.view_mut(),
3448 &typed_view("div", lhs)?,
3449 &typed_view("div", rhs)?,
3450 |x, y| x / y,
3451 )
3452 .map_err(|err| crate::Error::backend_failure("div", err))?;
3453 Ok(tensor_from_array(out))
3454 } else if lhs.shape().is_empty() {
3455 let scalar = typed_host_data("div", lhs)?[0];
3456 let mut out = unsafe { typed_array_uninit_from_pool(buffers, rhs.shape()) }?;
3458 map_into(&mut out.view_mut(), &typed_view("div", rhs)?, |x| {
3459 scalar / x
3460 })
3461 .map_err(|err| crate::Error::backend_failure("div", err))?;
3462 Ok(tensor_from_array(out))
3463 } else if rhs.shape().is_empty() {
3464 let scalar = typed_host_data("div", rhs)?[0];
3465 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3467 map_into(&mut out.view_mut(), &typed_view("div", lhs)?, |x| {
3468 x / scalar
3469 })
3470 .map_err(|err| crate::Error::backend_failure("div", err))?;
3471 Ok(tensor_from_array(out))
3472 } else {
3473 Err(crate::Error::ShapeMismatch {
3474 op: "div",
3475 lhs: lhs.shape().to_vec(),
3476 rhs: rhs.shape().to_vec(),
3477 })
3478 }
3479}
3480
3481fn typed_integer_div_with_pool<T>(
3482 buffers: &mut BufferPool,
3483 lhs: &TypedTensor<T>,
3484 rhs: &TypedTensor<T>,
3485) -> crate::Result<TypedTensor<T>>
3486where
3487 T: WrappingIntegerElem,
3488{
3489 let rhs_view = typed_view("div", rhs)?;
3490 ensure_no_zero_divisor("div", &rhs_view)?;
3491 typed_binary_with_pool("div", buffers, lhs, rhs, |x, y| x.wrapping_div_elem(y))
3492}
3493
3494fn typed_integer_div_view_with_pool<T, L, R>(
3495 buffers: &mut BufferPool,
3496 lhs: &TypedTensorView<'_, T, L>,
3497 rhs: &TypedTensorView<'_, T, R>,
3498) -> crate::Result<TypedTensor<T>>
3499where
3500 T: WrappingIntegerElem,
3501 L: TensorRank,
3502 R: TensorRank,
3503{
3504 let rhs_view = typed_view_from_view("div", rhs)?;
3505 ensure_no_zero_divisor("div", &rhs_view)?;
3506 typed_binary_view_with_pool("div", buffers, lhs, rhs, |x, y| x.wrapping_div_elem(y))
3507}
3508
3509fn typed_rem_with_pool<T>(
3510 buffers: &mut BufferPool,
3511 lhs: &TypedTensor<T>,
3512 rhs: &TypedTensor<T>,
3513) -> crate::Result<TypedTensor<T>>
3514where
3515 T: Copy + Clone + Zero + StdRem<Output = T> + PoolScalar + 'static,
3516{
3517 typed_binary_with_pool("rem", buffers, lhs, rhs, |x, y| x % y)
3518}
3519
3520fn typed_integer_rem_with_pool<T>(
3521 buffers: &mut BufferPool,
3522 lhs: &TypedTensor<T>,
3523 rhs: &TypedTensor<T>,
3524) -> crate::Result<TypedTensor<T>>
3525where
3526 T: WrappingIntegerElem,
3527{
3528 let rhs_view = typed_view("rem", rhs)?;
3529 ensure_no_zero_divisor("rem", &rhs_view)?;
3530 typed_binary_with_pool("rem", buffers, lhs, rhs, |x, y| x.wrapping_rem_elem(y))
3531}
3532
3533fn typed_integer_rem_view_with_pool<T, L, R>(
3534 buffers: &mut BufferPool,
3535 lhs: &TypedTensorView<'_, T, L>,
3536 rhs: &TypedTensorView<'_, T, R>,
3537) -> crate::Result<TypedTensor<T>>
3538where
3539 T: WrappingIntegerElem,
3540 L: TensorRank,
3541 R: TensorRank,
3542{
3543 let rhs_view = typed_view_from_view("rem", rhs)?;
3544 ensure_no_zero_divisor("rem", &rhs_view)?;
3545 typed_binary_view_with_pool("rem", buffers, lhs, rhs, |x, y| x.wrapping_rem_elem(y))
3546}
3547
3548pub(crate) fn typed_neg_with_pool<T>(
3549 buffers: &mut BufferPool,
3550 input: &TypedTensor<T>,
3551) -> crate::Result<TypedTensor<T>>
3552where
3553 T: Copy + Clone + Zero + Neg<Output = T> + PoolScalar,
3554{
3555 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
3557 map_into(&mut out.view_mut(), &typed_view("neg", input)?, |x| -x)
3558 .map_err(|err| crate::Error::backend_failure("neg", err))?;
3559 Ok(tensor_from_array(out))
3560}
3561
3562fn typed_wrapping_unary_with_pool<T>(
3563 op: &'static str,
3564 buffers: &mut BufferPool,
3565 input: &TypedTensor<T>,
3566 f: impl Fn(T) -> T + Copy + Sync,
3567) -> crate::Result<TypedTensor<T>>
3568where
3569 T: WrappingIntegerElem,
3570{
3571 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
3573 map_into(&mut out.view_mut(), &typed_view(op, input)?, f)
3574 .map_err(|err| crate::Error::backend_failure(op, err))?;
3575 Ok(tensor_from_array(out))
3576}
3577
3578fn typed_wrapping_neg_with_pool<T>(
3579 buffers: &mut BufferPool,
3580 input: &TypedTensor<T>,
3581) -> crate::Result<TypedTensor<T>>
3582where
3583 T: WrappingIntegerElem,
3584{
3585 typed_wrapping_unary_with_pool("neg", buffers, input, |x| x.wrapping_neg_elem())
3586}
3587
3588fn typed_wrapping_abs_with_pool<T>(
3589 buffers: &mut BufferPool,
3590 input: &TypedTensor<T>,
3591) -> crate::Result<TypedTensor<T>>
3592where
3593 T: WrappingIntegerElem,
3594{
3595 typed_wrapping_unary_with_pool("abs", buffers, input, |x| x.wrapping_abs_elem())
3596}
3597
3598pub(crate) fn typed_conj_with_pool<T>(
3599 buffers: &mut BufferPool,
3600 input: &TypedTensor<T>,
3601) -> crate::Result<TypedTensor<T>>
3602where
3603 T: Copy + Clone + Zero + ConjElem + PoolScalar,
3604{
3605 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
3607 map_into(&mut out.view_mut(), &typed_view("conj", input)?, |x| {
3608 x.conj_elem()
3609 })
3610 .map_err(|err| crate::Error::backend_failure("conj", err))?;
3611 Ok(tensor_from_array(out))
3612}
3613
3614pub(crate) fn typed_abs_with_pool<T>(
3615 buffers: &mut BufferPool,
3616 input: &TypedTensor<T>,
3617) -> crate::Result<TypedTensor<T>>
3618where
3619 T: Tier2Elem + PoolScalar,
3620{
3621 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
3623 map_into(&mut out.view_mut(), &typed_view("abs", input)?, |x| {
3624 x.abs_elem()
3625 })
3626 .map_err(|err| crate::Error::backend_failure("abs", err))?;
3627 Ok(tensor_from_array(out))
3628}
3629
3630fn typed_complex_abs_with_pool<T>(
3631 buffers: &mut BufferPool,
3632 input: &TypedTensor<Complex<T>>,
3633) -> crate::Result<TypedTensor<T>>
3634where
3635 T: num_traits::Float + PoolScalar,
3636{
3637 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
3639 map_into(&mut out.view_mut(), &typed_view("abs", input)?, |x| {
3640 x.norm()
3641 })
3642 .map_err(|err| crate::Error::backend_failure("abs", err))?;
3643 Ok(tensor_from_array(out))
3644}
3645
3646fn typed_complex_abs_view_with_pool<T, R>(
3647 buffers: &mut BufferPool,
3648 input: &TypedTensorView<'_, Complex<T>, R>,
3649) -> crate::Result<TypedTensor<T>>
3650where
3651 T: num_traits::Float + PoolScalar + 'static,
3652 R: TensorRank,
3653{
3654 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
3656 map_into(
3657 &mut out.view_mut(),
3658 &typed_view_from_view("abs", input)?,
3659 |x| x.norm(),
3660 )
3661 .map_err(|err| crate::Error::backend_failure("abs", err))?;
3662 Ok(tensor_from_array(out))
3663}
3664
3665pub(crate) fn typed_sign_with_pool<T>(
3666 buffers: &mut BufferPool,
3667 input: &TypedTensor<T>,
3668) -> crate::Result<TypedTensor<T>>
3669where
3670 T: Tier2Elem + PoolScalar,
3671{
3672 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
3674 map_into(&mut out.view_mut(), &typed_view("sign", input)?, |x| {
3675 x.sign_elem()
3676 })
3677 .map_err(|err| crate::Error::backend_failure("sign", err))?;
3678 Ok(tensor_from_array(out))
3679}
3680
3681fn typed_integer_sign_with_pool<T>(
3682 buffers: &mut BufferPool,
3683 input: &TypedTensor<T>,
3684) -> crate::Result<TypedTensor<T>>
3685where
3686 T: WrappingIntegerElem,
3687{
3688 typed_wrapping_unary_with_pool("sign", buffers, input, |x| x.signum_elem())
3689}
3690
3691pub(crate) fn typed_maximum_with_pool<T>(
3692 buffers: &mut BufferPool,
3693 lhs: &TypedTensor<T>,
3694 rhs: &TypedTensor<T>,
3695) -> crate::Result<TypedTensor<T>>
3696where
3697 T: OrderedElem + PoolScalar,
3698{
3699 if lhs.shape() != rhs.shape() {
3700 return Err(crate::Error::ShapeMismatch {
3701 op: "maximum",
3702 lhs: lhs.shape().to_vec(),
3703 rhs: rhs.shape().to_vec(),
3704 });
3705 }
3706 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3708 zip_map2_into(
3709 &mut out.view_mut(),
3710 &typed_view("maximum", lhs)?,
3711 &typed_view("maximum", rhs)?,
3712 |x, y| x.max_elem(y),
3713 )
3714 .map_err(|err| crate::Error::backend_failure("maximum", err))?;
3715 Ok(tensor_from_array(out))
3716}
3717
3718pub(crate) fn typed_minimum_with_pool<T>(
3719 buffers: &mut BufferPool,
3720 lhs: &TypedTensor<T>,
3721 rhs: &TypedTensor<T>,
3722) -> crate::Result<TypedTensor<T>>
3723where
3724 T: OrderedElem + PoolScalar,
3725{
3726 if lhs.shape() != rhs.shape() {
3727 return Err(crate::Error::ShapeMismatch {
3728 op: "minimum",
3729 lhs: lhs.shape().to_vec(),
3730 rhs: rhs.shape().to_vec(),
3731 });
3732 }
3733 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3735 zip_map2_into(
3736 &mut out.view_mut(),
3737 &typed_view("minimum", lhs)?,
3738 &typed_view("minimum", rhs)?,
3739 |x, y| x.min_elem(y),
3740 )
3741 .map_err(|err| crate::Error::backend_failure("minimum", err))?;
3742 Ok(tensor_from_array(out))
3743}
3744
3745pub(crate) fn typed_compare_with_pool<T>(
3746 buffers: &mut BufferPool,
3747 lhs: &TypedTensor<T>,
3748 rhs: &TypedTensor<T>,
3749 dir: &CompareDir,
3750) -> crate::Result<TypedTensor<bool>>
3751where
3752 T: CompareElem,
3753{
3754 if lhs.shape() != rhs.shape() {
3755 return Err(crate::Error::ShapeMismatch {
3756 op: "compare",
3757 lhs: lhs.shape().to_vec(),
3758 rhs: rhs.shape().to_vec(),
3759 });
3760 }
3761 let mut out = unsafe { typed_array_uninit_from_pool(buffers, lhs.shape()) }?;
3763 zip_map2_into(
3764 &mut out.view_mut(),
3765 &typed_view("compare", lhs)?,
3766 &typed_view("compare", rhs)?,
3767 |x, y| x.compare_elem(y, dir),
3768 )
3769 .map_err(|err| crate::Error::backend_failure("compare", err))?;
3770 Ok(tensor_from_array(out))
3771}
3772
3773pub(crate) fn typed_select_with_pool<T>(
3774 buffers: &mut BufferPool,
3775 pred: &TypedTensor<bool>,
3776 on_true: &TypedTensor<T>,
3777 on_false: &TypedTensor<T>,
3778) -> crate::Result<TypedTensor<T>>
3779where
3780 T: Copy + PoolScalar,
3781{
3782 if pred.shape() != on_true.shape() {
3783 return Err(crate::Error::ShapeMismatch {
3784 op: "select",
3785 lhs: pred.shape().to_vec(),
3786 rhs: on_true.shape().to_vec(),
3787 });
3788 }
3789 if pred.shape() != on_false.shape() {
3790 return Err(crate::Error::ShapeMismatch {
3791 op: "select",
3792 lhs: pred.shape().to_vec(),
3793 rhs: on_false.shape().to_vec(),
3794 });
3795 }
3796 let mut out = unsafe { typed_array_uninit_from_pool(buffers, pred.shape()) }?;
3798 zip_map3_into(
3799 &mut out.view_mut(),
3800 &typed_view("select", pred)?,
3801 &typed_view("select", on_true)?,
3802 &typed_view("select", on_false)?,
3803 |p, t, f| if p { t } else { f },
3804 )
3805 .map_err(|err| crate::Error::backend_failure("select", err))?;
3806 Ok(tensor_from_array(out))
3807}
3808
3809pub(crate) fn typed_clamp_with_pool<T>(
3810 buffers: &mut BufferPool,
3811 input: &TypedTensor<T>,
3812 lower: &TypedTensor<T>,
3813 upper: &TypedTensor<T>,
3814) -> crate::Result<TypedTensor<T>>
3815where
3816 T: OrderedElem + PoolScalar,
3817{
3818 if input.shape() != lower.shape() {
3819 return Err(crate::Error::ShapeMismatch {
3820 op: "clamp",
3821 lhs: input.shape().to_vec(),
3822 rhs: lower.shape().to_vec(),
3823 });
3824 }
3825 if input.shape() != upper.shape() {
3826 return Err(crate::Error::ShapeMismatch {
3827 op: "clamp",
3828 lhs: input.shape().to_vec(),
3829 rhs: upper.shape().to_vec(),
3830 });
3831 }
3832 let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
3834 zip_map3_into(
3835 &mut out.view_mut(),
3836 &typed_view("clamp", input)?,
3837 &typed_view("clamp", lower)?,
3838 &typed_view("clamp", upper)?,
3839 |x, lo, hi| hi.min_elem(lo.max_elem(x)),
3840 )
3841 .map_err(|err| crate::Error::backend_failure("clamp", err))?;
3842 Ok(tensor_from_array(out))
3843}
3844
3845#[cfg(test)]
3846mod tests;