svod_tensor/bitwise.rs
1//! Bitwise operations for integer tensors.
2//!
3//! This module provides bitwise operations: AND, OR, XOR, left shift, right shift.
4//! All operations require integer or boolean dtypes.
5
6use snafu::ResultExt;
7
8use super::*;
9
10impl Tensor {
11 /// Bitwise AND operation.
12 ///
13 /// Performs element-wise bitwise AND between two tensors with broadcasting.
14 /// Both tensors must have integer or boolean dtype.
15 #[track_caller]
16 pub fn bitwise_and(&self, other: &Tensor) -> Result<Tensor> {
17 let (lhs, rhs) = self.broadcast_for_binop(other)?;
18 lhs.uop().try_and_op(&rhs.uop()).map(Self::new).context(UOpSnafu)
19 }
20
21 /// Bitwise OR operation.
22 ///
23 /// Performs element-wise bitwise OR between two tensors with broadcasting.
24 /// Both tensors must have integer or boolean dtype.
25 #[track_caller]
26 pub fn bitwise_or(&self, other: &Tensor) -> Result<Tensor> {
27 let (lhs, rhs) = self.broadcast_for_binop(other)?;
28 lhs.uop().try_or_op(&rhs.uop()).map(Self::new).context(UOpSnafu)
29 }
30
31 /// Bitwise XOR operation.
32 ///
33 /// Performs element-wise bitwise XOR between two tensors with broadcasting.
34 /// Both tensors must have integer or boolean dtype.
35 #[track_caller]
36 pub fn bitwise_xor(&self, other: &Tensor) -> Result<Tensor> {
37 let (lhs, rhs) = self.broadcast_for_binop(other)?;
38 lhs.uop().try_xor_op(&rhs.uop()).map(Self::new).context(UOpSnafu)
39 }
40
41 /// Left shift operation.
42 ///
43 /// Shifts bits of the tensor to the left by the specified amount with broadcasting.
44 /// The tensor must have integer or boolean dtype.
45 #[track_caller]
46 pub fn lshift(&self, other: &Tensor) -> Result<Tensor> {
47 let (lhs, rhs) = self.broadcast_for_binop(other)?;
48 lhs.uop().try_shl_op(&rhs.uop()).map(Self::new).context(UOpSnafu)
49 }
50
51 /// Right shift operation.
52 ///
53 /// Shifts bits of the tensor to the right by the specified amount with broadcasting.
54 /// The tensor must have integer or boolean dtype.
55 #[track_caller]
56 pub fn rshift(&self, other: &Tensor) -> Result<Tensor> {
57 let (lhs, rhs) = self.broadcast_for_binop(other)?;
58 lhs.uop().try_shr_op(&rhs.uop()).map(Self::new).context(UOpSnafu)
59 }
60}