diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index e1d9155..5680fef 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/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, ) diff --git a/tests/test_unified_visitor_regressions.py b/tests/test_unified_visitor_regressions.py index 9ae3691..9c30e4a 100644 --- a/tests/test_unified_visitor_regressions.py +++ b/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]))