AI fix and add fp8 support

more_math/Parser/UnifiedMathVisitor.py

@@ -2181,15 +2181,69 @@ class UnifiedMathVisitor(MathExprVisitor):
         v = (yield ctx.expr())
         return self._bitwise_not(v)
 
+    def _bitwise_op(self, a, b, torch_op, scalar_op):
+        """Binary bitwise operation handler supporting tensors, lists, and scalars."""
+        if self._is_tensor(a) and a.numel() == 1:
+            a = int(a.flatten()[0].item())
+        if self._is_tensor(b) and b.numel() == 1:
+            b = int(b.flatten()[0].item())
+
+        # Handle tensor-list combinations
+        if self._is_tensor(a) and self._is_list(b):
+            if a.shape[0] == len(b):
+                A = torch.split(a, 1)
+                return torch.cat([self._bitwise_op(x, y, torch_op, scalar_op) for x, y in zip(A, b)], dim=0)
+            return torch.cat([self._bitwise_op(a, x, torch_op, scalar_op) for x in b], dim=0)
+        if self._is_list(a) and self._is_tensor(b):
+            if b.shape[0] == len(a):
+                B = torch.split(b, 1)
+                return torch.cat([self._bitwise_op(x, y, torch_op, scalar_op) for x, y in zip(a, B)], dim=0)
+            return torch.cat([self._bitwise_op(x, b, torch_op, scalar_op) for x in a], dim=0)
+
+        # Handle list-list and list-scalar combinations
+        if self._is_list(a) and not self._is_tensor(b):
+            if self._is_list(b):
+                if len(a) != len(b):
+                    raise ValueError("List length mismatch in bitwise operation")
+                return [self._bitwise_op(x, y, torch_op, scalar_op) for x, y in zip(a, b)]
+            return [self._bitwise_op(x, b, torch_op, scalar_op) for x in a]
+
+        if not self._is_tensor(a) and self._is_list(b):
+            return [self._bitwise_op(a, x, torch_op, scalar_op) for x in b]
+
+        # Handle tensor operations
+        if self._is_tensor(a) or self._is_tensor(b):
+            if torch_op:
+                return torch_op(a, b).contiguous()
+            return scalar_op(a, b)
+
+        return scalar_op(a, b)
+
+    def _get_bitwise_view_dtype(self, elem_size):
+        """Get appropriate integer dtype for bitwise operations based on element size."""
+        if elem_size == 1:
+            return torch.int8
+        elif elem_size == 2:
+            return torch.int16
+        elif elem_size == 4:
+            return torch.int32
+        elif elem_size == 8:
+            return torch.int64
+        else:
+            # Fallback for unusual sizes
+            return torch.int32
+
     def _bitwise_not(self, v):
-        """Unary bitwise NOT handling for tensors, lists and scalars."""
+        """Unary bitwise NOT handling for tensors, lists and scalars with support for fp16 and int16."""
         if self._is_tensor(v):
             t = self._promote_to_tensor(v)
             elem_size = t.element_size() if hasattr(t, 'element_size') else 4
-            view_dtype = torch.int64 if elem_size == 8 else torch.int32
+            view_dtype = self._get_bitwise_view_dtype(elem_size)
+            original_dtype = t.dtype
+            
             bits = t.view(view_dtype)
             res_bits = torch.bitwise_not(bits)
-            return res_bits.view(t.dtype)
+            return res_bits.view(original_dtype).contiguous()
 
         if self._is_list(v):
             return [self._bitwise_not(x) for x in v]