hpt_traits/ops/binary.rs
1use std::borrow::BorrowMut;
2
3use hpt_common::error::base::TensorError;
4use hpt_types::dtype::TypeCommon;
5
6use crate::tensor::CommonBounds;
7
8/// A trait for binary operations on tensors.
9pub trait NormalBinOps<RHS = Self>
10where
11 <<Self as NormalBinOps<RHS>>::OutputMeta as TypeCommon>::Vec: Send + Sync,
12{
13 /// The output tensor type.
14 type Output;
15 /// The output tensor data type.
16 type OutputMeta: CommonBounds;
17 /// The inplace output tensor type.
18 type InplaceOutput;
19
20 /// add with specified output tensor
21 ///
22 /// ## Parameters:
23 /// `rhs`: The right-hand side tensor.
24 ///
25 /// `out`: The output tensor.
26 ///
27 /// ## Example:
28 /// ```rust
29 /// let a = Tensor::<f32>::new([2.0]);
30 /// let b = Tensor::<f32>::new([3.0]);
31 /// let c = a.add_(&b, &mut a.clone())?; // c and a point to the same memory
32 /// println!("{}", c); // [5.0]
33 /// ```
34 fn add_<U>(&self, rhs: RHS, out: U) -> std::result::Result<Self::Output, TensorError>
35 where
36 U: BorrowMut<Self::InplaceOutput>;
37
38 /// subtract with specified output tensor
39 ///
40 /// ## Parameters:
41 /// `rhs`: The right-hand side tensor.
42 ///
43 /// `out`: The output tensor.
44 ///
45 /// ## Example:
46 /// ```rust
47 /// let a = Tensor::<f32>::new([2.0]);
48 /// let b = Tensor::<f32>::new([3.0]);
49 /// let c = a.sub_(&b, &mut a.clone())?; // c and a point to the same memory
50 /// println!("{}", c); // [-1.0]
51 /// ```
52 fn sub_<U>(&self, rhs: RHS, out: U) -> std::result::Result<Self::Output, TensorError>
53 where
54 U: BorrowMut<Self::InplaceOutput>;
55
56 /// multiply with specified output tensor
57 ///
58 /// ## Parameters:
59 /// `rhs`: The right-hand side tensor.
60 ///
61 /// `out`: The output tensor.
62 ///
63 /// ## Example:
64 /// ```rust
65 /// let a = Tensor::<f32>::new([2.0]);
66 /// let b = Tensor::<f32>::new([3.0]);
67 /// let c = a.mul_(&b, &mut a.clone())?; // c and a point to the same memory
68 /// println!("{}", c); // [6.0]
69 /// ```
70 fn mul_<U>(&self, rhs: RHS, out: U) -> std::result::Result<Self::Output, TensorError>
71 where
72 U: BorrowMut<Self::InplaceOutput>;
73
74 /// rem with specified output tensor
75 ///
76 /// ## Parameters:
77 /// `rhs`: The right-hand side tensor.
78 ///
79 /// `out`: The output tensor.
80 ///
81 /// ## Example:
82 /// ```rust
83 /// let a = Tensor::<f32>::new([2.0]);
84 /// let b = Tensor::<f32>::new([3.0]);
85 /// let c = a.mul_(&b, &mut a.clone())?; // c and a point to the same memory
86 /// println!("{}", c); // [6.0]
87 /// ```
88 fn rem_<U>(&self, rhs: RHS, out: U) -> std::result::Result<Self::Output, TensorError>
89 where
90 U: BorrowMut<Self::InplaceOutput>;
91}
92
93/// A trait for binary operations on tensors.
94pub trait FloatBinOps<RHS = Self> {
95 /// The output tensor type.
96 type Output;
97 /// The output tensor data type.
98 type OutputMeta: CommonBounds;
99 /// The inplace output tensor type.
100 type InplaceOutput;
101
102 /// Compute `sqrt(x^2 + y^2)` for all elements
103 ///
104 /// ## Parameters:
105 /// `rhs`: The right-hand side tensor.
106 ///
107 /// ## Example:
108 /// ```rust
109 /// let base = Tensor::<f32>::new(&[3.0, 4.0, 5.0]);
110 /// let height = Tensor::<f32>::new(&[4.0, 3.0, 12.0]);
111 /// let hypotenuse = base.hypot(&height)?; // [5.0, 5.0, 13.0]
112 /// let fixed_height = base.hypot(4.0)?; // [5.0, 5.66, 6.40]
113 /// ```
114 fn hypot<B>(&self, rhs: B) -> std::result::Result<Self::Output, TensorError>
115 where
116 B: Into<RHS>;
117
118 /// hypot with specified output tensor
119 ///
120 /// ## Parameters:
121 /// `rhs`: The right-hand side tensor.
122 ///
123 /// `out`: The output tensor.
124 ///
125 /// ## Example:
126 /// ```rust
127 /// let base = Tensor::<f32>::new(&[3.0, 4.0, 5.0]);
128 /// let height = Tensor::<f32>::new(&[4.0, 3.0, 12.0]);
129 /// let result = base.hypot_(&height, &mut base.clone())?; // [5.0, 5.0, 13.0]
130 /// ```
131 fn hypot_<B, U>(&self, rhs: B, out: U) -> std::result::Result<Self::Output, TensorError>
132 where
133 U: BorrowMut<Self::InplaceOutput>,
134 B: Into<RHS>;
135
136 /// division with specified output tensor
137 ///
138 /// ## Parameters:
139 /// `rhs`: The right-hand side tensor.
140 ///
141 /// `out`: The output tensor.
142 ///
143 /// ## Example:
144 /// ```rust
145 /// let a = Tensor::<f32>::new([2.0]);
146 /// let b = Tensor::<f32>::new([3.0]);
147 /// let c = a.div_(&b, &mut a.clone())?; // [0.6667]
148 /// ```
149 fn div_<B, U>(&self, rhs: B, out: U) -> std::result::Result<Self::Output, TensorError>
150 where
151 U: BorrowMut<Self::InplaceOutput>,
152 B: Into<RHS>;
153
154 /// Power of `self` and `rhs`
155 ///
156 /// ## Parameters:
157 /// `rhs`: The right-hand side tensor.
158 ///
159 /// ## Example:
160 /// ```rust
161 /// let a = Tensor::<f32>::new(&[2.0, 3.0, 4.0]);
162 /// let b = Tensor::<f32>::new(&[2.0, 3.0, 2.0]);
163 /// let c = a.pow(&b)?; // [4.0, 27.0, 16.0]
164 /// let d = a.pow(2.0f64)?; // [4.0, 9.0, 16.0]
165 /// ```
166 fn pow<B>(&self, rhs: B) -> std::result::Result<Self::Output, TensorError>
167 where
168 B: Into<RHS>;
169
170 /// Power of `self` and `rhs` with specified output tensor
171 ///
172 /// ## Parameters:
173 /// `rhs`: The right-hand side tensor.
174 ///
175 /// `out`: The output tensor.
176 ///
177 /// ## Example:
178 /// ```rust
179 /// let a = Tensor::<f32>::new(&[2.0, 3.0, 4.0]);
180 /// let b = Tensor::<f32>::new(&[2.0, 3.0, 2.0]);
181 /// let c = a.pow_(&b, &mut a.clone())?; // [4.0, 27.0, 16.0]
182 /// ```
183 fn pow_<B, U>(&self, rhs: B, out: U) -> std::result::Result<Self::Output, TensorError>
184 where
185 U: BorrowMut<Self::InplaceOutput>,
186 B: Into<RHS>;
187}
188
189/// A trait for matrix multiplication operations on tensors.
190pub trait Matmul<RHS = Self>
191where
192 <<Self as Matmul<RHS>>::OutputMeta as TypeCommon>::Vec: Send + Sync,
193{
194 /// The output tensor type.
195 type Output;
196 /// The output tensor data type.
197 type OutputMeta: CommonBounds;
198 /// The inplace output tensor type.
199 type InplaceOutput;
200
201 /// Perform matrix multiplication of two tensors. The behavior depends on the dimensions of the input tensors:
202 ///
203 /// - If both tensors are 2D, they are multiplied as matrices
204 /// - If either tensor is ND (N > 2), it is treated as a stack of matrices
205 /// - Broadcasting is applied to match dimensions
206 ///
207 /// ## Parameters:
208 /// `rhs`: The right-hand side tensor.
209 ///
210 /// ## Example:
211 /// ```rust
212 /// // 2D matrix multiplication
213 /// let a = Tensor::<f64>::new(&[[1., 2.], [3., 4.]]);
214 /// let b = Tensor::<f64>::new(&[[5., 6.], [7., 8.]]);
215 /// let c = a.matmul(&b)?;
216 /// println!("2D result:\n{}", c);
217 ///
218 /// // 3D batch matrix multiplication
219 /// let d = Tensor::<f64>::ones(&[2, 2, 3])?; // 2 matrices of shape 2x3
220 /// let e = Tensor::<f64>::ones(&[2, 3, 2])?; // 2 matrices of shape 3x2
221 /// let f = d.matmul(&e)?; // 2 matrices of shape 2x2
222 /// println!("3D result:\n{}", f);
223 /// ```
224 #[track_caller]
225 fn matmul(&self, rhs: RHS) -> std::result::Result<Self::Output, TensorError>;
226
227 /// matrix multiplication with specified output tensor
228 ///
229 /// ## Parameters:
230 /// `rhs`: The right-hand side tensor.
231 ///
232 /// `out`: The output tensor.
233 ///
234 /// ## Example:
235 /// ```rust
236 /// // 2D matrix multiplication
237 /// let a = Tensor::<f64>::new(&[[1., 2.], [3., 4.]]);
238 /// let b = Tensor::<f64>::new(&[[5., 6.], [7., 8.]]);
239 /// let c = a.matmul_(&b, &mut a.clone())?;
240 /// println!("2D result:\n{}", c);
241 ///
242 /// // 3D batch matrix multiplication
243 /// let d = Tensor::<f64>::ones(&[2, 2, 3])?; // 2 matrices of shape 2x3
244 /// let e = Tensor::<f64>::ones(&[2, 3, 2])?; // 2 matrices of shape 3x2
245 /// let f = d.matmul_(&e, &mut d.clone())?; // 2 matrices of shape 2x2
246 /// println!("3D result:\n{}", f);
247 /// ```
248 #[track_caller]
249 fn matmul_<U>(&self, rhs: RHS, out: U) -> std::result::Result<Self::InplaceOutput, TensorError>
250 where
251 U: BorrowMut<Self::InplaceOutput> + BorrowMut<Self::InplaceOutput>;
252}
253
254/// A trait for matrix multiplication operations on tensors with post-operation.
255pub trait MatmulPost<RHS = Self>
256where
257 <<Self as MatmulPost<RHS>>::OutputMeta as TypeCommon>::Vec: Send + Sync,
258{
259 /// The output tensor type.
260 type Output;
261 /// The output tensor data type.
262 type OutputMeta: CommonBounds;
263 /// The inplace output tensor type.
264 type InplaceOutput;
265
266 /// Same as `matmul` but will perform post operation before writing final result to the memory.
267 ///
268 /// ## Parameters:
269 /// `rhs`: The right-hand side tensor.
270 ///
271 /// `post_op`: post operation function that takes scalar `T` as input and returns scalar `T`
272 ///
273 /// `post_op_vec`: post operation function that takes simd vector `T::Simd` as input and returns simd vector `T::Simd`
274 ///
275 /// ## Example:
276 /// ```rust
277 /// // 2D matrix multiplication
278 /// let a = Tensor::<f64>::new(&[[1., 2.], [3., 4.]]);
279 /// let b = Tensor::<f64>::new(&[[5., 6.], [7., 8.]]);
280 /// let c = a.matmul_post(&b, |x| x._relu(), |x| x._relu())?;
281 /// println!("2D result:\n{}", c);
282 ///
283 /// // 3D batch matrix multiplication
284 /// let d = Tensor::<f64>::ones(&[2, 2, 3])?; // 2 matrices of shape 2x3
285 /// let e = Tensor::<f64>::ones(&[2, 3, 2])?; // 2 matrices of shape 3x2
286 /// let f = d.matmul_post(&e, |x| x._relu(), |x| x._relu())?; // 2 matrices of shape 2x2
287 /// println!("3D result:\n{}", f);
288 /// ```
289 #[track_caller]
290 fn matmul_post(
291 &self,
292 rhs: RHS,
293 post_op: fn(Self::OutputMeta) -> Self::OutputMeta,
294 post_op_vec: fn(
295 <<Self as MatmulPost<RHS>>::OutputMeta as TypeCommon>::Vec,
296 ) -> <<Self as MatmulPost<RHS>>::OutputMeta as TypeCommon>::Vec,
297 ) -> std::result::Result<Self::Output, TensorError>;
298
299 /// matrix multiplication with specified output tensor and post operation
300 ///
301 /// ## Parameters:
302 /// `rhs`: The right-hand side tensor.
303 ///
304 /// `out`: The output tensor.
305 ///
306 /// `post_op`: post operation function that takes scalar `T` as input and returns scalar `T`
307 ///
308 /// `post_op_vec`: post operation function that takes simd vector `T::Simd` as input and returns simd vector `T::Simd`
309 ///
310 /// ## Example:
311 /// ```rust
312 /// // 2D matrix multiplication
313 /// let a = Tensor::<f64>::new(&[[1., 2.], [3., 4.]]);
314 /// let b = Tensor::<f64>::new(&[[5., 6.], [7., 8.]]);
315 /// let c = a.matmul_post_(&b, &mut a.clone(), |x| x._relu(), |x| x._relu())?;
316 /// println!("2D result:\n{}", c);
317 ///
318 /// // 3D batch matrix multiplication
319 /// let d = Tensor::<f64>::ones(&[2, 2, 3])?; // 2 matrices of shape 2x3
320 /// let e = Tensor::<f64>::ones(&[2, 3, 2])?; // 2 matrices of shape 3x2
321 /// let f = d.matmul_post_(&e, &mut d.clone(), |x| x._relu(), |x| x._relu())?; // 2 matrices of shape 2x2
322 /// println!("3D result:\n{}", f);
323 /// ```
324 #[track_caller]
325 fn matmul_post_<U>(
326 &self,
327 rhs: RHS,
328 post_op: fn(Self::OutputMeta) -> Self::OutputMeta,
329 post_op_vec: fn(
330 <<Self as MatmulPost<RHS>>::OutputMeta as TypeCommon>::Vec,
331 ) -> <<Self as MatmulPost<RHS>>::OutputMeta as TypeCommon>::Vec,
332 out: U,
333 ) -> std::result::Result<Self::InplaceOutput, TensorError>
334 where
335 U: BorrowMut<Self::InplaceOutput> + BorrowMut<Self::InplaceOutput>;
336}
337/// A trait for gemm operations on tensors.
338pub trait Gemm<RHS = Self>
339where
340 <<Self as Gemm<RHS>>::OutputMeta as TypeCommon>::Vec: Send + Sync,
341{
342 /// The output tensor type.
343 type Output;
344 /// The output tensor data type.
345 type OutputMeta: CommonBounds;
346 /// The inplace output tensor type.
347 type InplaceOutput;
348
349 /// Perform gemm (general matrix multiplication) of two tensors. The behavior depends on the dimensions of the input tensors:
350 ///
351 /// - If both tensors are 2D, they are multiplied as matrices
352 /// - If either tensor is ND (N > 2), it is treated as a stack of matrices
353 /// - Broadcasting is applied to match dimensions
354 ///
355 /// ## Parameters:
356 /// `rhs`: The right-hand side tensor.
357 ///
358 /// `alpha`: Scaling factor for the matrix product (A @ B)
359 ///
360 /// `beta`: Scaling factor for the existing values in output matrix C
361 ///
362 /// `conj_dst`: Whether to conjugate C before scaling with beta
363 ///
364 /// `conj_lhs`: Whether to conjugate A before multiplication
365 ///
366 /// `conj_rhs`: Whether to conjugate B before multiplication
367 ///
368 /// ## Example:
369 /// ```rust
370 /// // 2D matrix multiplication
371 /// let a = Tensor::<f64>::new(&[[1., 2.], [3., 4.]]);
372 /// let b = Tensor::<f64>::new(&[[5., 6.], [7., 8.]]);
373 /// let c = a.gemm(&b, 0.0, 1.0, false, false, false)?;
374 /// println!("2D result:\n{}", c);
375 ///
376 /// // 3D batch matrix multiplication
377 /// let d = Tensor::<f64>::ones(&[2, 2, 3])?; // 2 matrices of shape 2x3
378 /// let e = Tensor::<f64>::ones(&[2, 3, 2])?; // 2 matrices of shape 3x2
379 /// let f = d.gemm(&e, 0.0, 1.0, false, false, false)?; // 2 matrices of shape 2x2
380 /// println!("3D result:\n{}", f);
381 /// ```
382 #[track_caller]
383 fn gemm(
384 &self,
385 rhs: RHS,
386 alpha: Self::OutputMeta,
387 beta: Self::OutputMeta,
388 conj_dst: bool,
389 conj_lhs: bool,
390 conj_rhs: bool,
391 ) -> std::result::Result<Self::Output, TensorError>;
392
393 /// gemm (general matrix multiplication) with specified output tensor
394 ///
395 /// ## Parameters:
396 /// `rhs`: The right-hand side tensor.
397 ///
398 /// `alpha`: Scaling factor for the matrix product (A @ B)
399 ///
400 /// `beta`: Scaling factor for the existing values in output matrix C
401 ///
402 /// `conj_dst`: Whether to conjugate C before scaling with beta
403 ///
404 /// `conj_lhs`: Whether to conjugate A before multiplication
405 ///
406 /// `conj_rhs`: Whether to conjugate B before multiplication
407 ///
408 /// `out`: The output tensor.
409 ///
410 /// ## Example:
411 /// ```rust
412 /// // 2D matrix multiplication
413 /// let a = Tensor::<f64>::new(&[[1., 2.], [3., 4.]]);
414 /// let b = Tensor::<f64>::new(&[[5., 6.], [7., 8.]]);
415 /// let c = a.gemm_(&b, 0.0, 1.0, false, false, false, &mut a.clone())?;
416 /// println!("2D result:\n{}", c);
417 ///
418 /// // 3D batch matrix multiplication
419 /// let d = Tensor::<f64>::ones(&[2, 2, 3])?; // 2 matrices of shape 2x3
420 /// let e = Tensor::<f64>::ones(&[2, 3, 2])?; // 2 matrices of shape 3x2
421 /// let f = d.gemm_(&e, 0.0, 1.0, false, false, false, &mut d.clone())?; // 2 matrices of shape 2x2
422 /// println!("3D result:\n{}", f);
423 /// ```
424 #[track_caller]
425 fn gemm_<U>(
426 &self,
427 rhs: RHS,
428 alpha: Self::OutputMeta,
429 beta: Self::OutputMeta,
430 conj_dst: bool,
431 conj_lhs: bool,
432 conj_rhs: bool,
433 out: U,
434 ) -> std::result::Result<Self::InplaceOutput, TensorError>
435 where
436 U: BorrowMut<Self::InplaceOutput> + BorrowMut<Self::InplaceOutput>;
437}
438
439/// A trait for tensor dot operations on tensors.
440pub trait TensorDot<RHS = Self> {
441 /// The output tensor type.
442 type Output;
443
444 /// Compute tensor dot product along specified axes. This is a generalization of matrix multiplication to higher dimensions.
445 ///
446 /// ## Parameters:
447 /// `rhs`: The right-hand side tensor.
448 ///
449 /// `axes`: A tuple of two arrays specifying the axes to contract over:
450 /// - First array contains axes from the first tensor
451 /// - Second array contains axes from the second tensor
452 /// - Arrays must have same length N
453 ///
454 /// ## Example:
455 /// ```rust
456 /// // Matrix multiplication (2D tensordot)
457 /// let a = Tensor::new(&[[1., 2.], [3., 4.]]);
458 /// let b = Tensor::new(&[[5., 6.], [7., 8.]]);
459 /// let c = a.tensordot(&b, ([1], [0]))?; // Contract last axis of a with first axis of b
460 /// println!("Matrix multiplication:\n{}", c);
461 ///
462 /// // Higher dimensional example
463 /// let d = Tensor::<f32>::ones(&[2, 3, 4])?;
464 /// let e = Tensor::<f32>::ones(&[4, 3, 2])?;
465 /// let f = d.tensordot(&e, ([1, 2], [1, 0]))?; // Contract axes 1,2 of d with axes 1,0 of e
466 /// println!("Higher dimensional result:\n{}", f);
467 /// ```
468 fn tensordot<const N: usize>(
469 &self,
470 rhs: &RHS,
471 axes: ([i64; N], [i64; N]),
472 ) -> std::result::Result<Self::Output, TensorError>;
473}