Skip to main content

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}