logical and bitwise functions

heat/core/arithmetics.py

@@ -4,10 +4,12 @@
 from . import dndarray
 from . import operations
 from . import stride_tricks
+from . import types
 
 __all__ = [
     "add",
     "bitwise_and",
+    "bitwise_not",
     "bitwise_or",
     "bitwise_xor",
     "diff",
@@ -16,13 +18,16 @@
     "floordiv",
     "floor_divide",
     "fmod",
+    "invert",
+    "left_shift",
     "mod",
     "mul",
     "multiply",
     "pow",
-    "prod",
     "power",
+    "prod",
     "remainder",
+    "right_shift",
     "sub",
     "subtract",
     "sum",
@@ -100,7 +105,13 @@ def bitwise_and(t1, t2):
     >>> ht.bitwise_and(ht.array([True, True]), ht.array([False, True]))
     tensor([False,  True])
     """
-    return operations.__binary_bit_op("__and__", t1, t2)
+    dtypes = (types.heat_type_of(t1), types.heat_type_of(t2))
+
+    for dtype in dtypes:
+        if types.heat_type_is_inexact(dtype):
+            raise TypeError("Operation is not supported for float types")
+
+    return operations.__binary_op(torch.Tensor.__and__, t1, t2)
 
 
 def bitwise_or(t1, t2):
@@ -138,7 +149,13 @@ def bitwise_or(t1, t2):
     >>> ht.bitwise_or(ht.array([True, True]), ht.array([False, True]))
     tensor([ True,  True])
     """
-    return operations.__binary_bit_op("__or__", t1, t2)
+    dtypes = (types.heat_type_of(t1), types.heat_type_of(t2))
+
+    for dtype in dtypes:
+        if types.heat_type_is_inexact(dtype):
+            raise TypeError("Operation is not supported for float types")
+
+    return operations.__binary_op(torch.Tensor.__or__, t1, t2)
 
 
 def bitwise_xor(t1, t2):
@@ -171,7 +188,13 @@ def bitwise_xor(t1, t2):
     >>> ht.bitwise_xor(ht.array([True, True]), ht.array([False, True]))
     tensor([ True, False])
     """
-    return operations.__binary_bit_op("__xor__", t1, t2)
+    dtypes = (types.heat_type_of(t1), types.heat_type_of(t2))
+
+    for dtype in dtypes:
+        if types.heat_type_is_inexact(dtype):
+            raise TypeError("Operation is not supported for float types")
+
+    return operations.__binary_op(torch.Tensor.__xor__, t1, t2)
 
 
 def diff(a, n=1, axis=-1):
@@ -380,6 +403,65 @@ def floordiv(t1, t2):
 floor_divide = floordiv
 
 
+def invert(t, out=None):
+    """
+    Computes the bitwise NOT of the given input tensor. The input tensor must be of integral or Boolean types. For bool tensors, it computes the logical NOT.
+    Bitwise_not is an alias for invert.
+
+    Returns
+        -------
+        result: ht.DNDarray
+            A tensor containing the results of element-wise inversion.
+
+    Examples:
+    ---------
+    >>> ht.invert(ht.array([13], dtype=ht.uint8))
+    tensor([242], dtype=ht.uint8)
+    >>> ht.bitwise_not(ht.array([-1, -2, 3], dtype=ht.int8))
+    tensor([ 0,  1, -4], dtype=ht.int8)
+    """
+    dtype = types.heat_type_of(t)
+
+    if types.heat_type_is_inexact(dtype):
+        raise TypeError("Operation is not supported for float types")
+
+    return operations.__local_op(torch.bitwise_not, t, out, no_cast=True)
+
+
+# alias for invert
+bitwise_not = invert
+
+
+def left_shift(t1, t2):
+    """
+    Shift the bits of an integer to the left.
+
+    Parameters
+    ----------
+    t1: scalar or tensor
+
+    t2: scalar or tensor
+        integer number of zero bits to add
+
+    Returns
+    -------
+    result: ht.NDNarray
+        A tensor containing the results of element-wise left shift operation.
+
+    Examples:
+    ---------
+    >>> ht.left_shift(ht.array[1,2,3], 1)
+    tensor([2, 4, 6])
+    """
+    dtypes = (types.heat_type_of(t1), types.heat_type_of(t2))
+
+    for dtype in dtypes:
+        if not types.heat_type_is_exact(dtype):
+            raise TypeError("Operation is supported for integer types only")
+
+    return operations.__binary_op(torch.Tensor.__lshift__, t1, t2)
+
+
 def mod(t1, t2):
     """
     Element-wise division remainder of values of operand t1 by values of operand t2 (i.e. t1 % t2), not commutative.
@@ -546,6 +628,36 @@ def remainder(t1, t2):
     return operations.__binary_op(torch.remainder, t1, t2)
 
 
+def right_shift(t1, t2):
+    """
+    Shift the bits of an integer to the right.
+
+    Parameters
+    ----------
+    t1: scalar or tensor
+
+    t2: scalar or tensor
+        integer number of bits to remove
+
+    Returns
+    -------
+    result: ht.NDNarray
+        A tensor containing the results of element-wise right shift operation.
+
+    Examples:
+    ---------
+    >>> ht.right_shift(ht.array[1,2,3], 1)
+    tensor([0, 1, 1])
+    """
+    dtypes = (types.heat_type_of(t1), types.heat_type_of(t2))
+
+    for dtype in dtypes:
+        if not types.heat_type_is_exact(dtype):
+            raise TypeError("Operation is supported for integer types only")
+
+    return operations.__binary_op(torch.Tensor.__rshift__, t1, t2)
+
+
 def prod(x, axis=None, out=None, keepdim=None):
     """
     Return the product of array elements over a given axis.

heat/core/dndarray.py

@@ -1531,6 +1531,17 @@ def __int__(self):
         """
         return self.__cast(int)
 
+    def __invert__(self):
+        """
+        Bit-wise inversion, or bit-wise NOT, element-wise.
+
+        Returns
+        -------
+        result: ht.DNDarray
+            A tensor containing the results of element-wise inversion.
+        """
+        return arithmetics.invert(self)
+
     def is_balanced(self):
         """
         Determine if a DNDarray is balanced evenly (or as evenly as possible) across all nodes
@@ -1702,6 +1713,27 @@ def log1p(self, out=None):
         """
         return exponential.log1p(self, out)
 
+    def __lshift__(self, other):
+        """
+        Shift the bits of an integer to the left.
+
+        Parameters
+        ----------
+        other: scalar or tensor
+           number of zero bits to add
+
+        Returns
+        -------
+        result: ht.NDNarray
+           A tensor containing the results of element-wise left shift operation.
+
+        Examples:
+        ---------
+        >>> ht.array([1, 2, 4]) << 1
+        tensor([2, 4, 8])
+        """
+        return arithmetics.left_shift(self, other)
+
     def __lt__(self, other):
         """
         Element-wise rich comparison of relation "less than" with values from second operand (scalar or tensor)
@@ -2473,6 +2505,27 @@ def __rpow__(self, other):
         """
         return arithmetics.pow(other, self)
 
+    def __rshift__(self, other):
+        """
+        Shift the bits of an integer to the right.
+
+        Parameters
+        ----------
+        other: scalar or tensor
+           number of bits to remove
+
+        Returns
+        -------
+        result: ht.NDNarray
+           A tensor containing the results of element-wise right shift operation.
+
+        Examples:
+        ---------
+        >>> ht.array([1, 2, 4]) >> 1
+        tensor([0, 1, 2])
+        """
+        return arithmetics.right_shift(self, other)
+
     def __rsub__(self, other):
         """
         Element-wise subtraction of another tensor or a scalar from the tensor.

heat/core/logical.py

@@ -6,8 +6,9 @@
 from . import manipulations
 from . import operations
 from . import dndarray
+from . import types
 
-__all__ = ["all", "allclose", "any"]
+__all__ = ["all", "allclose", "any", "logical_and", "logical_not", "logical_or", "logical_xor"]
 
 
 def all(x, axis=None, out=None, keepdim=None):
@@ -207,3 +208,98 @@ def local_any(t, *args, **kwargs):
         return torch.any(t != 0, *args, **kwargs)
 
     return operations.__reduce_op(x, local_any, MPI.LOR, axis=axis, out=out, keepdim=keepdim)
+
+
+def logical_and(t1, t2):
+    """
+    Compute the truth value of t1 AND t2 element-wise.
+
+    Parameters:
+    -----------
+    t1, t2: tensor
+        input tensors of same shape
+
+    Returns:
+    --------
+    boolean_tensor : tensor of type bool
+        Element-wise result of t1 AND t2.
+
+    Examples:
+    ---------
+    >>> ht.logical_and(ht.array([True, False]), ht.array([False, False]))
+    tensor([ False, False])
+    """
+    return operations.__binary_op(
+        torch.Tensor.__and__, types.bool(t1, device=t1.device), types.bool(t2, device=t2.device)
+    )
+
+
+def logical_not(t, out=None):
+    """
+    Computes the element-wise logical NOT of the given input tensor.
+
+    Parameters:
+    -----------
+    t1: tensor
+        input tensor
+    out : tensor, optional
+        Alternative output tensor in which to place the result. It must have the same shape as the expected output.
+        The output is a tensor with dtype=bool.
+
+    Returns:
+    --------
+    boolean_tensor : tensor of type bool
+        Element-wise result of NOT t.
+
+    Examples:
+    ---------
+    >>> ht.logical_not(ht.array([True, False]))
+    tensor([ False,  True])
+    """
+    return operations.__local_op(torch.logical_not, t, out)
+
+
+def logical_or(t1, t2):
+    """
+    Compute the truth value of t1 OR t2 element-wise.
+
+    Parameters:
+    -----------
+    t1, t2: tensor
+        input tensors of same shape
+
+    Returns:
+    --------
+    boolean_tensor : tensor of type bool
+        Element-wise result of t1 OR t2.
+
+    Examples:
+    ---------
+    >>> ht.logical_or(ht.array([True, False]), ht.array([False, False]))
+    tensor([True, False])
+    """
+    return operations.__binary_op(
+        torch.Tensor.__or__, types.bool(t1, device=t1.device), types.bool(t2, device=t2.device)
+    )
+
+
+def logical_xor(t1, t2):
+    """
+    Computes the element-wise logical XOR of the given input tensors.
+
+    Parameters:
+    -----------
+    t1, t2: tensor
+        input tensors of same shape
+
+    Returns:
+    --------
+    boolean_tensor : tensor of type bool
+        Element-wise result of t1 XOR t2.
+
+    Examples:
+    ---------
+    >>> ht.logical_xor(ht.array([True, False, True]), ht.array([True, False, False]))
+    tensor([ False, False,  True])
+    """
+    return operations.__binary_op(torch.logical_xor, t1, t2)