fix test

more_math/Parser/UnifiedMathVisitor.py

@@ -1146,11 +1146,22 @@ class UnifiedMathVisitor(MathExprVisitor):
     def visitHypotFunc(self, ctx):
         return self._bin_op((yield ctx.expr(0)), (yield ctx.expr(1)), torch.hypot, math.hypot, ctx)
 
+    def _logical_binary_tensor_op(self, a, b, operation):
+        if isinstance(a, torch.Tensor):
+            device = a.device
+        elif isinstance(b, torch.Tensor):
+            device = b.device
+        else:
+            return 1.0 if operation(bool(a), bool(b)) else 0.0
+        a = torch.as_tensor(a, device=device)
+        b = torch.as_tensor(b, device=device)
+        return operation(a, b).to(torch.float32)
+
     def visitLogicalAndFunc(self, ctx):
         return self._bin_op(
             (yield ctx.expr(0)),
             (yield ctx.expr(1)),
-            lambda a, b: torch.logical_and(a, b).to(torch.float32),
+            lambda a, b: self._logical_binary_tensor_op(a, b, torch.logical_and),
             lambda a, b: 1.0 if bool(a) and bool(b) else 0.0,
             ctx,
         )
@@ -1159,7 +1170,7 @@ class UnifiedMathVisitor(MathExprVisitor):
         return self._bin_op(
             (yield ctx.expr(0)),
             (yield ctx.expr(1)),
-            lambda a, b: torch.logical_or(a, b).to(torch.float32),
+            lambda a, b: self._logical_binary_tensor_op(a, b, torch.logical_or),
             lambda a, b: 1.0 if bool(a) or bool(b) else 0.0,
             ctx,
         )

tests/test_unified_visitor_regressions.py

@@ -50,13 +50,22 @@ def test_float8_operations_preserve_dtype(dtype_name):
     values = torch.tensor([1.0, 2.0], dtype=dtype)
     tensor_result = _visit("x * 2", {"x": values})
     unary_result = _visit("-x", {"x": values})
+    math_result = _visit("log2(x)", {"x": values})
+    hypot_result = _visit("hypot(3*x, 4*x)", {"x": values})
+    logical_result = _visit("logical_and(x, 1)", {"x": values})
     dict_result = _visit("V0 * 2", {"V0": ChangeTrackingDict({"weight": values})})
 
     assert tensor_result.dtype == dtype
     assert unary_result.dtype == dtype
+    assert math_result.dtype == dtype
+    assert hypot_result.dtype == dtype
+    assert logical_result.dtype == dtype
     assert dict_result["weight"].dtype == dtype
     assert torch.equal(tensor_result.float(), torch.tensor([2.0, 4.0]))
     assert torch.equal(unary_result.float(), torch.tensor([-1.0, -2.0]))
+    assert torch.equal(math_result.float(), torch.tensor([0.0, 1.0]))
+    assert torch.equal(hypot_result.float(), torch.tensor([5.0, 10.0]))
+    assert torch.equal(logical_result.float(), torch.tensor([1.0, 1.0]))
     assert torch.equal(dict_result["weight"].float(), torch.tensor([2.0, 4.0]))