diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index 7c3786d..c9b3480 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/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]