AI modernise tests

more_math/Parser/UnifiedMathVisitor.py

@@ -608,6 +608,12 @@ class UnifiedMathVisitor(MathExprVisitor):
 
         # Handle slice expressions - all cases with colons
         exprs = slice_elem_ctx.expr() if hasattr(slice_elem_ctx, 'expr') else []
+        if exprs is None:
+            exprs = []
+        elif isinstance(exprs, tuple):
+            exprs = list(exprs)
+        elif not isinstance(exprs, list):
+            exprs = [exprs]
 
         start = None
         stop = None
@@ -4817,6 +4823,10 @@ class UnifiedMathVisitor(MathExprVisitor):
             for v in val:
                 t.append(self.i2rgb(v))
             return t
+        if self._is_tensor(val):
+            val = val.to(torch.int64)
+        else:
+            val = int(val)
         r = (val >> 16) & 0xFF
         g = (val >> 8) & 0xFF
         b = val & 0xFF

tests/reproduce_indexing.py

@@ -7,9 +7,10 @@ import os
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
+from more_math.Parser.antlr_router import get_antlr_modules
 import antlr4
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
 
 class TestIndexing(unittest.TestCase):
     def evaluate(self, expr, variables):

tests/reproduce_size_error.py

@@ -9,10 +9,11 @@ if _project_root not in sys.path:
     sys.path.insert(0, _project_root)
 
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from antlr4 import InputStream, CommonTokenStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 def parse_and_visit(expr_str, variables):
     lexer = MathExprLexer(InputStream(expr_str))
     stream = CommonTokenStream(lexer)

tests/test_bitwise_comprehensive.py

@@ -30,9 +30,8 @@ def test_mixed_dtype_operations():
     a_int8 = torch.tensor([15, 7, 3], dtype=torch.int8)
     visitor = UnifiedMathVisitor({"a": a_int8})
     result = visitor._bitwise_not(a_int8)
-    print(f"   Input:  {a_int8} (dtype={a_int8.dtype})")
-    print(f"   ~Input: {result} (dtype={result.dtype})")
     assert result.dtype == torch.int8
+    assert torch.equal(result, torch.bitwise_not(a_int8))
 
     # int16 operations
     print("\n2. INT16 Bitwise Operations:")
@@ -50,24 +49,26 @@ def test_mixed_dtype_operations():
     print(f"   a | b: {result_or}")
     print(f"   a ^ b: {result_xor}")
     assert all(t.dtype == torch.int16 for t in [result_and, result_or, result_xor])
+    assert torch.equal(result_and, torch.tensor([15, 31, 7], dtype=torch.int16))
+    assert torch.equal(result_or, torch.tensor([255, 127, 63], dtype=torch.int16))
+    assert torch.equal(result_xor, torch.tensor([240, 96, 56], dtype=torch.int16))
 
     # fp16 operations
     print("\n3. FP16 Bitwise Operations (bit-level manipulation):")
     a_fp16 = torch.tensor([1.0, -2.0, 3.5], dtype=torch.float16)
     visitor = UnifiedMathVisitor({"a": a_fp16})
     result = visitor._bitwise_not(a_fp16)
-    print(f"   Input:         {a_fp16} (dtype={a_fp16.dtype})")
-    print(f"   Bit-flipped:   {result} (dtype={result.dtype})")
     assert result.dtype == torch.float16
+    expected = torch.bitwise_not(a_fp16.view(torch.int16)).view(torch.float16)
+    assert torch.equal(result, expected)
 
     # int32 operations (existing, should still work)
     print("\n4. INT32 Bitwise Operations (backward compatibility):")
     a_int32 = torch.tensor([65535, 32767, 16383], dtype=torch.int32)
     visitor = UnifiedMathVisitor({"a": a_int32})
     result = visitor._bitwise_not(a_int32)
-    print(f"   Input:  {a_int32} (dtype={a_int32.dtype})")
-    print(f"   ~Input: {result} (dtype={result.dtype})")
     assert result.dtype == torch.int32
+    assert torch.equal(result, torch.bitwise_not(a_int32))
 
     print("\n" + "=" * 70)
     print("✓ All mixed dtype operations completed successfully!")
@@ -95,28 +96,28 @@ def test_scalar_list_tensor_combinations():
     a = torch.tensor([7, 14, 21], dtype=torch.int16)
     b = torch.tensor([3, 5, 7], dtype=torch.int16)
     result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y, ctx)
-    print(f"   {a} & {b} = {result}")
+    assert torch.equal(result, torch.tensor([3, 4, 5], dtype=torch.int16))
 
     # Tensor & List
     print("\n2. Tensor & List (mixed):")
     a = torch.tensor([15, 14, 13], dtype=torch.int16)
     b = [7, 3, 1]
     result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y, ctx)
-    print(f"   Tensor({a}) & List({b}) = {result}")
+    assert torch.equal(result, torch.tensor([7, 2, 1], dtype=torch.int16))
 
     # Scalar & List
     print("\n3. Scalar & List:")
     a = 15
     b = [7, 3, 1]
     result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y, ctx)
-    print(f"   {a} & {b} = {result}")
+    assert result == [7, 3, 1]
 
     # List & List
     print("\n4. List & List:")
     a = [15, 14, 13]
     b = [7, 3, 1]
     result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y, ctx)
-    print(f"   {a} & {b} = {result}")
+    assert result == [7, 2, 1]
 
     print("\n" + "=" * 70)
     print("✓ All scalar/list/tensor combinations work correctly!")
@@ -153,12 +154,17 @@ def test_element_size_mapping():
         elif elem_size == 2:
             expected_view = "int16"
         elif elem_size == 4:
-            expected_view = "int32"
+            expected_view = torch.int32
         else:
-            expected_view = "int64"
+            expected_view = torch.int64
+        if elem_size == 1:
+            expected_view = torch.int8
+        elif elem_size == 2:
+            expected_view = torch.int16
 
         print(f"  {name:12} -> {elem_size} byte(s) -> {str(view_dtype):18}")
         assert elem_size == expected_elem_size
+        assert view_dtype == expected_view
 
     print("\n" + "=" * 70)
     print("✓ Element size mapping is correct!")

tests/test_bitwise_fix.py

@@ -1,133 +0,0 @@
-#!/usr/bin/env python3
-"""Quick test to verify bitwise operations work with tensors"""
-import sys
-sys.path.insert(0, r'D:\stability\Data\Packages\ComfyUI')
-
-import torch
-from custom_nodes.more_math.more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-
-def test_bitwise_xor_float():
-    """Test XOR with float tensors (the reported bug)"""
-    print("Testing bitwise XOR with float tensors...")
-
-    # Create float tensors (this was causing the error)
-    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
-    b = torch.tensor([4.0, 5.0, 6.0], dtype=torch.float32)
-
-    visitor = UnifiedMathVisitor({"a": a, "b": b})
-
-    try:
-        # This should convert to int64, perform XOR, then convert back
-        result = visitor._bitwise_op(a, b, torch.bitwise_xor, lambda x, y: x ^ y)
-        print("✓ XOR succeeded!")
-        print(f"  Input a (float32): {a}")
-        print(f"  Input b (float32): {b}")
-        print(f"  Result: {result}")
-        print(f"  Result dtype: {result.dtype}")
-        return True
-    except Exception as e:
-        print(f"✗ XOR failed: {e}")
-        return False
-
-def test_bitwise_and_float():
-    """Test AND with float tensors"""
-    print("\nTesting bitwise AND with float tensors...")
-
-    a = torch.tensor([15.0, 14.0, 13.0], dtype=torch.float32)
-    b = torch.tensor([7.0, 3.0, 1.0], dtype=torch.float32)
-
-    visitor = UnifiedMathVisitor({"a": a, "b": b})
-
-    try:
-        result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y)
-        print("✓ AND succeeded!")
-        print(f"  Input a (float32): {a}")
-        print(f"  Input b (float32): {b}")
-        print(f"  Result: {result}")
-        print(f"  Result dtype: {result.dtype}")
-        return True
-    except Exception as e:
-        print(f"✗ AND failed: {e}")
-        return False
-
-def test_bitwise_or_float():
-    """Test OR with float tensors"""
-    print("\nTesting bitwise OR with float tensors...")
-
-    a = torch.tensor([15.0, 14.0, 13.0], dtype=torch.float32)
-    b = torch.tensor([7.0, 3.0, 1.0], dtype=torch.float32)
-
-    visitor = UnifiedMathVisitor({"a": a, "b": b})
-
-    try:
-        result = visitor._bitwise_op(a, b, torch.bitwise_or, lambda x, y: x | y)
-        print("✓ OR succeeded!")
-        print(f"  Input a (float32): {a}")
-        print(f"  Input b (float32): {b}")
-        print(f"  Result: {result}")
-        print(f"  Result dtype: {result.dtype}")
-        return True
-    except Exception as e:
-        print(f"✗ OR failed: {e}")
-        return False
-
-def test_bitwise_not_float():
-    """Test NOT with float tensors"""
-    print("\nTesting bitwise NOT with float tensors...")
-
-    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
-
-    visitor = UnifiedMathVisitor({"a": a})
-
-    try:
-        result = visitor._bitwise_not(a)
-        print("✓ NOT succeeded!")
-        print(f"  Input a (float32): {a}")
-        print(f"  Result: {result}")
-        print(f"  Result dtype: {result.dtype}")
-        return True
-    except Exception as e:
-        print(f"✗ NOT failed: {e}")
-        return False
-
-def test_int_tensors_preserved():
-    """Ensure int tensors still work as before"""
-    print("\nTesting that int tensor dtypes are preserved...")
-
-    a = torch.tensor([15, 14, 13], dtype=torch.int16)
-    b = torch.tensor([7, 3, 1], dtype=torch.int16)
-
-    visitor = UnifiedMathVisitor({"a": a, "b": b})
-
-    try:
-        result = visitor._bitwise_op(a, b, torch.bitwise_and, lambda x, y: x & y)
-        assert result.dtype == torch.int16, f"Expected int16, got {result.dtype}"
-        print("✓ Int16 dtype preserved!")
-        print(f"  Input a (int16): {a}")
-        print(f"  Input b (int16): {b}")
-        print(f"  Result (int16): {result}")
-        return True
-    except Exception as e:
-        print(f"✗ Int16 test failed: {e}")
-        return False
-
-if __name__ == "__main__":
-    print("=" * 70)
-    print("Bitwise Operations Fix Verification")
-    print("=" * 70)
-
-    results = [
-        test_bitwise_xor_float(),
-        test_bitwise_and_float(),
-        test_bitwise_or_float(),
-        test_bitwise_not_float(),
-        test_int_tensors_preserved(),
-    ]
-
-    print("\n" + "=" * 70)
-    if all(results):
-        print("✓ All tests passed!")
-    else:
-        print("✗ Some tests failed")
-        sys.exit(1)
-    print("=" * 70)

tests/test_bitwise_shifts.py

@@ -6,9 +6,11 @@ import torch
 import sys
 sys.path.insert(0, 'custom_nodes/more_math')
 
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from antlr4 import InputStream, CommonTokenFactory
+from antlr4 import InputStream, CommonTokenStream
+
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
 
 def parse_and_evaluate(expression, variables=None):
     """Parse and evaluate a math expression"""
@@ -16,9 +18,9 @@ def parse_and_evaluate(expression, variables=None):
         variables = {}
 
     input_stream = InputStream(expression)
-    lexer = MathExprParser(input_stream).lexer
-    stream = CommonTokenFactory()
-    parser = MathExprParser(input_stream)
+    lexer = MathExprLexer(input_stream)
+    stream = CommonTokenStream(lexer)
+    parser = MathExprParser(stream)
     tree = parser.start()
 
     visitor = UnifiedMathVisitor(variables, device='cpu')
@@ -26,137 +28,66 @@ def parse_and_evaluate(expression, variables=None):
     return result
 
 def test_bit_shifts():
-    """Test bitwise shift operators"""
-    print("=" * 60)
-    print("Testing Bitwise Shift Operators")
-    print("=" * 60)
-
+    """Test scalar and chained integer shifts."""
     test_cases = [
-        # Left shift: 5 << 2 = 20 (0101 << 2 = 10100)
-        ("5 << 2", {}, 20),
-
-        # Right shift: 20 >> 2 = 5 (10100 >> 2 = 0101)
-        ("20 >> 2", {}, 5),
-
-        # Left shift with variable
+        ("a << b", {"a": 5, "b": 2}, 20),
+        ("a >> b", {"a": 20, "b": 2}, 5),
         ("x << 3", {"x": 4}, 32),  # 4 << 3 = 32
-
-        # Right shift with variable
         ("x >> 2", {"x": 16}, 4),  # 16 >> 2 = 4
-
-        # Chained shifts
-        ("(8 << 2) >> 3", {}, 4),  # (32) >> 3 = 4
+        ("(a << b) >> c", {"a": 8, "b": 2, "c": 3}, 4),
     ]
 
     for expr, vars, expected in test_cases:
-        try:
-            result = parse_and_evaluate(expr, vars)
-            status = "✓" if result == expected else "✗"
-            print(f"{status} {expr:30} = {result:10} (expected {expected})")
-        except Exception as e:
-            print(f"✗ {expr:30} ERROR: {e}")
+        assert parse_and_evaluate(expr, vars) == expected
 
 def test_bit_count():
-    """Test bit count function"""
-    print("\n" + "=" * 60)
-    print("Testing Bitwise Bit Count Function")
-    print("=" * 60)
-
+    """Test scalar bit counts."""
     test_cases = [
-        # bitcount(5) = 2  (0101 has 2 set bits)
         ("bitcount(5)", {}, 2),
-
-        # bitcount(15) = 4 (1111 has 4 set bits)
         ("bitcount(15)", {}, 4),
-
-        # bitcount(7) = 3 (111 has 3 set bits)
         ("bitcount(7)", {}, 3),
-
-        # bitcount(255) = 8 (11111111 has 8 set bits)
         ("bitcount(255)", {}, 8),
-
-        # bitcount(0) = 0
         ("bitcount(0)", {}, 0),
-
-        # With variable
         ("bitcount(x)", {"x": 31}, 5),  # 31 = 11111 = 5 bits set
     ]
 
     for expr, vars, expected in test_cases:
-        try:
-            result = parse_and_evaluate(expr, vars)
-            status = "✓" if result == expected else "✗"
-            print(f"{status} {expr:30} = {result:10} (expected {expected})")
-        except Exception as e:
-            print(f"✗ {expr:30} ERROR: {e}")
+        assert parse_and_evaluate(expr, vars) == expected
 
 def test_bit_shifts_with_tensors():
-    """Test bitwise shifts with tensors"""
-    print("\n" + "=" * 60)
-    print("Testing Bitwise Shifts with Tensors")
-    print("=" * 60)
-
+    """Test bitwise shifts with integer tensors."""
     vars = {
         "a": torch.tensor([1, 2, 4, 8], dtype=torch.int32),
         "shift": 2,
     }
 
-    try:
-        result = parse_and_evaluate("a << shift", vars)
-        expected = torch.tensor([4, 8, 16, 32], dtype=torch.int32)
-        match = torch.equal(result, expected)
-        status = "✓" if match else "✗"
-        print(f"{status} tensor_shift_left: [1,2,4,8] << 2 = {result.tolist()}")
-    except Exception as e:
-        print(f"✗ tensor_shift_left ERROR: {e}")
+    result = parse_and_evaluate("a << shift", vars)
+    expected = torch.tensor([4, 8, 16, 32], dtype=torch.int32)
+    assert torch.equal(result, expected)
 
-    try:
-        result = parse_and_evaluate("a >> shift", vars)
-        expected = torch.tensor([0, 0, 1, 2], dtype=torch.int32)
-        match = torch.equal(result, expected)
-        status = "✓" if match else "✗"
-        print(f"{status} tensor_shift_right: [1,2,4,8] >> 2 = {result.tolist()}")
-    except Exception as e:
-        print(f"✗ tensor_shift_right ERROR: {e}")
+    result = parse_and_evaluate("a >> shift", vars)
+    expected = torch.tensor([0, 0, 1, 2], dtype=torch.int32)
+    assert torch.equal(result, expected)
 
 def test_bit_count_with_tensors():
-    """Test bit count with tensors"""
-    print("\n" + "=" * 60)
-    print("Testing Bit Count with Tensors")
-    print("=" * 60)
-
+    """Test bit counts with integer tensors."""
     vars = {
         "nums": torch.tensor([5, 15, 7, 255], dtype=torch.int32),
     }
 
-    try:
-        result = parse_and_evaluate("bitcount(nums)", vars)
-        # Expected: [2, 4, 3, 8] set bits
-        print(f"✓ bitcount([5,15,7,255]): {result}")
-    except Exception as e:
-        print(f"✗ bitcount_tensor ERROR: {e}")
+    result = parse_and_evaluate("bitcount(nums)", vars)
+    expected = torch.tensor([2, 4, 3, 8], dtype=torch.float32)
+    assert torch.equal(result, expected)
 
 def test_combinations():
-    """Test combinations of shift and bit count"""
-    print("\n" + "=" * 60)
-    print("Testing Combinations")
-    print("=" * 60)
-
+    """Test combined shifts and bitwise functions."""
     test_cases = [
-        # Shift then count bits
-        ("bitcount(5 << 2)", {}, 2),  # 5 << 2 = 20 (10100) = 2 bits
-
-        # Combined with other operators
-        ("(5 << 2) | (3 << 4)", {}, 0xCC),  # 20 | 48 = 0xCC = 204
+        ("bitcount(a << b)", {"a": 5, "b": 2}, 2),
+        ("bor(a << b, c << d)", {"a": 5, "b": 2, "c": 3, "d": 4}, 52),
     ]
 
     for expr, vars, expected in test_cases:
-        try:
-            result = parse_and_evaluate(expr, vars)
-            status = "✓" if result == expected else "✗"
-            print(f"{status} {expr:35} = {result:10} (expected {expected})")
-        except Exception as e:
-            print(f"✗ {expr:35} ERROR: {e}")
+        assert parse_and_evaluate(expr, vars) == expected
 
 if __name__ == "__main__":
     test_bit_shifts()
@@ -165,6 +96,3 @@ if __name__ == "__main__":
     test_bit_count_with_tensors()
     test_combinations()
 
-    print("\n" + "=" * 60)
-    print("All tests completed!")
-    print("=" * 60)

tests/test_convolution.py

@@ -10,10 +10,11 @@ if _project_root not in sys.path:
     sys.path.insert(0, _project_root)
 
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from antlr4 import InputStream, CommonTokenStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 def parse_and_visit(expr_str, variables):
     lexer = MathExprLexer(InputStream(expr_str))
     stream = CommonTokenStream(lexer)

tests/test_error_messages.py

@@ -1,5 +1,6 @@
 import sys
 import os
+import pytest
 
 # Ensure we can import the module
 _here = os.path.abspath(os.path.dirname(__file__))
@@ -8,10 +9,11 @@ if _project_root not in sys.path:
     sys.path.insert(0, _project_root)
 
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from antlr4 import InputStream, CommonTokenStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 def parse_and_visit(expr_str, variables):
     lexer = MathExprLexer(InputStream(expr_str))
     stream = CommonTokenStream(lexer)
@@ -30,13 +32,10 @@ def test_variable_not_found_error():
         raise AssertionError("Expected ValueError not raised")
 
 def test_unknown_function_error():
-    try:
+    with pytest.raises(ValueError) as exc_info:
         parse_and_visit("unknown_func(1)", {})
-    except ValueError as e:
-        msg = str(e)
-        assert "1:0: Unknown function: unknown_func" in msg
-    else:
-        raise AssertionError("Expected ValueError not raised")
+    msg = str(exc_info.value)
+    assert "1:0: Unknown function or variable: unknown_func" in msg
 
 def test_get_value_invalid_arg_error():
     try:

tests/test_guider_math.py

@@ -28,7 +28,7 @@ class TestMathGuider(unittest.TestCase):
         # Expression: Average V0 and V1
         expr = "V0 * 0.5 + V1 * 0.5"
 
-        math_guider = MathGuider(V, F, expr, expr)  # Add expression1 parameter
+        math_guider = MathGuider(V, F, expr, expr, {})
 
         # Pseudo input
         x = torch.zeros((1, 4, 16, 16))
@@ -48,7 +48,7 @@ class TestMathGuider(unittest.TestCase):
         # a = V0, w = F0
         expr = "a + w"
 
-        math_guider = MathGuider(V, F, expr, expr)  # Add expression1 parameter
+        math_guider = MathGuider(V, F, expr, expr, {})
         x = torch.zeros((1, 4, 8, 8))
         sigma = torch.tensor(1.0)
 
@@ -60,7 +60,7 @@ class TestMathGuider(unittest.TestCase):
         g0 = MockGuider(1.0)
         V = {"V0": g0}
         F = {}
-        math_guider = MathGuider(V, F, "V0", "V0")  # Add expression1 parameter
+        math_guider = MathGuider(V, F, "V0", "V0", {})
 
         # Check if the property exists and matches g0's patcher
         self.assertIsNotNone(math_guider.model_patcher)
@@ -72,7 +72,7 @@ class TestMathGuider(unittest.TestCase):
         del g0.model_patcher # force remove
         V = {"V0": g0}
         F = {}
-        math_guider = MathGuider(V, F, "V0", "V0")  # Add expression1 parameter
+        math_guider = MathGuider(V, F, "V0", "V0", {})
         self.assertIsNone(math_guider.model_patcher)
 
     def test_math_guider_steps_context(self):
@@ -82,7 +82,7 @@ class TestMathGuider(unittest.TestCase):
         V = {"V0": g0}
 
         expr = "steps > 0 ? current_step / steps : 0.0"  # Handle division by zero
-        math_guider = MathGuider(V, {}, expr, expr)
+        math_guider = MathGuider(V, {}, expr, expr, {})
         math_guider.sigmas = sigmas # sets sigmas directly for testing
         math_guider.steps = len(sigmas) - 1  # steps = num_sigmas - 1 (2 steps)
 
@@ -112,7 +112,7 @@ class TestMathGuider(unittest.TestCase):
         # Input shape is (1, 4, 4, 4) = (batch, channel, height, width)
         # crop(a, [0, 1, 1, 1], [1, 2, 2, 2]) extracts 2x2x2 region
         expr = "crop(a, [0, 1, 1, 1], [1, 2, 2, 2])"
-        math_guider = MathGuider(V, F, expr, expr)
+        math_guider = MathGuider(V, F, expr, expr, {})
 
         # Input: 4D tensor (batch, channel, height, width)
         x = torch.ones((1, 4, 4, 4))
@@ -131,7 +131,7 @@ class TestMathGuider(unittest.TestCase):
 
         # Simple numeric expression
         expr = "a * 2"
-        math_guider = MathGuider(V, F, expr, expr)
+        math_guider = MathGuider(V, F, expr, expr, {})
 
         x = torch.ones((1, 4, 4, 4))
         sigma = torch.tensor(1.0)
@@ -150,7 +150,7 @@ class TestMathGuider(unittest.TestCase):
 
         # crop with 4D parameters for 4D input (batch, channel, height, width)
         expr = "crop(a, [0, 0, 0, 0], [1, 2, 2, 2])"
-        math_guider = MathGuider(V, F, expr, expr)
+        math_guider = MathGuider(V, F, expr, expr, {})
 
         x = torch.ones((1, 4, 4, 4))
         sigma = torch.tensor(1.0)

tests/test_model_math.py

@@ -49,7 +49,7 @@ def test_model_math_simple_add():
     expr = "a + b"
 
     # Execute
-    result_tuple = ModelMathNode.execute(Expression=expr, V={"V0": a, "V1": b}, F={})
+    result_tuple = ModelMathNode.execute(Expression=expr, V={"V0": a, "V1": b}, F={}, use_compute_device=False)
     result_model = result_tuple[0]
 
     # Verify patches
@@ -81,7 +81,7 @@ def test_model_math_zero_default():
     # diff = 10 - 10 = 0.
     # Should skip patch.
 
-    result_tuple = ModelMathNode.execute(Expression="a + b", V={"V0": a}, F={})
+    result_tuple = ModelMathNode.execute(Expression="a + b", V={"V0": a}, F={}, use_compute_device=False)
     result_model = result_tuple[0]
 
     assert len(result_model.patches) == 0
@@ -94,7 +94,7 @@ def test_model_math_custom_formula():
     # expr: a * w
     # w input = 0.5
 
-    result_tuple = ModelMathNode.execute(Expression="a * w", V={"V0": a}, F={"F0": 0.5})
+    result_tuple = ModelMathNode.execute(Expression="a * w", V={"V0": a}, F={"F0": 0.5}, use_compute_device=False)
     result_model = result_tuple[0]
 
     # res = 2.0 * 0.5 = 1.0
@@ -113,8 +113,8 @@ def test_model_math_device_mismatch():
     # But we can verify that the scalar created by visitor.visitNumberExp has the same device as 'a'.
 
     from antlr4 import InputStream, CommonTokenStream
-    from more_math.Parser.MathExprLexer import MathExprLexer
-    from more_math.Parser.MathExprParser import MathExprParser
+    from more_math.Parser.antlr_router import get_antlr_modules
+    MathExprLexer, MathExprParser, _ = get_antlr_modules()
 
     # Mock a tensor with a mocked .device property if we can't use real devices
     # But torch tensors created with torch.tensor are real.
@@ -168,7 +168,7 @@ def test_clip_math():
     sd_a = {"text_model.encoder.layers.0.mlp.fc1.weight": torch.tensor([1.0])}
     clip_a = MockPatcherContainer(sd_a)
 
-    result_tuple = CLIPMathNode.execute(Expression="a * 2", V={"V0": clip_a}, F={})
+    result_tuple = CLIPMathNode.execute(Expression="a * 2", V={"V0": clip_a}, F={}, use_compute_device=False)
     result_clip = result_tuple[0]
 
     # 1.0 * 2 = 2.0. diff = 1.0.
@@ -193,7 +193,7 @@ def test_vae_math():
 
     vae_a = MockVAE(sd_a)
 
-    result_tuple = VAEMathNode.execute(Expression="a + 1", V={"V0": vae_a}, F={})
+    result_tuple = VAEMathNode.execute(Expression="a + 1", V={"V0": vae_a}, F={}, use_compute_device=False)
     result_vae = result_tuple[0]
 
     # Check if result is a copy (not same object)
@@ -224,7 +224,7 @@ def test_model_math_disjoint_keys():
     # common: 1 + 2 = 3. diff = 2.
     # extra_b: should NOT be in patches because loop is over 'a'
 
-    result_tuple = ModelMathNode.execute(Expression="a + b", V={"V0": a.patcher, "V1": b.patcher}, F={})
+    result_tuple = ModelMathNode.execute(Expression="a + b", V={"V0": a.patcher, "V1": b.patcher}, F={}, use_compute_device=False)
     patches = result_tuple[0].patches
 
     assert "common" in patches
@@ -246,7 +246,7 @@ def test_model_math_weighted_merge():
     # expr: lerp(a, b, w) with w=0.5 -> 15.0
     # diff = 15 - 10 = 5.
 
-    result_tuple = ModelMathNode.execute(Expression="lerp(a, b, w)", V={"V0": a.patcher, "V1": b.patcher}, F={"F0": 0.5})
+    result_tuple = ModelMathNode.execute(Expression="lerp(a, b, w)", V={"V0": a.patcher, "V1": b.patcher}, F={"F0": 0.5}, use_compute_device=False)
     patches = result_tuple[0].patches
 
     assert "w" in patches
@@ -267,7 +267,7 @@ def test_model_math_3_way_merge():
     # Expected: 1 + 2 + 3 = 6.
     # Diff vs a (1.0) = 5.0.
 
-    result_tuple = ModelMathNode.execute(Expression="a + b + c", V={"V0": a.patcher, "V1": b.patcher, "V2": c.patcher}, F={})
+    result_tuple = ModelMathNode.execute(Expression="a + b + c", V={"V0": a.patcher, "V1": b.patcher, "V2": c.patcher}, F={}, use_compute_device=False)
     patches = result_tuple[0].patches
 
     assert "w" in patches

tests/test_more_math.py

@@ -42,7 +42,7 @@ def test_conditioning_math_node_initialization():
 def test_conditioning_math_node_metadata():
     assert ConditioningMathNode.RETURN_TYPES == ["CONDITIONING", "STACK"]
     assert ConditioningMathNode.FUNCTION == "EXECUTE_NORMALIZED"
-    assert ConditioningMathNode.CATEGORY == "More math"
+    assert ConditioningMathNode.CATEGORY == "More math/math"
 
 
 def test_latent_math_node_initialization():
@@ -53,7 +53,7 @@ def test_latent_math_node_initialization():
 def test_latent_math_node_metadata():
     assert LatentMathNode.RETURN_TYPES == ["LATENT", "STACK"]
     assert LatentMathNode.FUNCTION == "EXECUTE_NORMALIZED"
-    assert LatentMathNode.CATEGORY == "More math"
+    assert LatentMathNode.CATEGORY == "More math/math"
 
 
 def test_image_math_node_initialization():
@@ -64,7 +64,7 @@ def test_image_math_node_initialization():
 def test_image_math_node_metadata():
     assert ImageMathNode.RETURN_TYPES == ["IMAGE", "STACK"]
     assert ImageMathNode.FUNCTION == "EXECUTE_NORMALIZED"
-    assert ImageMathNode.CATEGORY == "More math"
+    assert ImageMathNode.CATEGORY == "More math/math"
 
 
 # ==========================================

tests/test_permute_fix.py

@@ -1,51 +1,17 @@
 import torch
-import sys
-import os
 
-# Add paths
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))) # ComfyUI root
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "more_math", "Parser")))
-
-import optical_flow_utils as ofu
-from comfy.nested_tensor import NestedTensor
+from more_math.Parser import optical_flow_utils as ofu
+from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
+from more_math.helper_functions import parse_expr
 
 def test_permute_and_apply():
-    print("\n--- Testing Permute and Flow Apply ---")
+    image = torch.arange(24.0).reshape(1, 2, 4, 3)
+    tree = parse_expr("permute(image, [0, 3, 1, 2])")
+    result = UnifiedMathVisitor({"image": image}).visit(tree)
+    assert torch.equal(result, image.permute(0, 3, 1, 2))
 
-    # 1. Test Permute on NestedTensor
-    print("Testing NestedTensor.permute...")
-    nt = NestedTensor([torch.randn(1, 256, 256, 3) for _ in range(2)])
-    try:
-        # Reorder [B, H, W, C] to [B, C, H, W]
-        nt_permuted = nt.permute(0, 3, 1, 2)
-        print(f"Original shape: {nt.shape}")
-        print(f"Permuted shape: {nt_permuted.shape}")
-        if nt_permuted.shape == (1, 3, 256, 256):
-            print("NestedTensor.permute SUCCESS!")
-        else:
-            print(f"NestedTensor.permute FAILED: Got {nt_permuted.shape}")
-    except Exception as e:
-        print(f"NestedTensor.permute FAILED with error: {e}")
-        import traceback
-        traceback.print_exc()
-
-    # 2. Test apply_flow
-    print("\nTesting ofu.apply_flow...")
-    img = torch.randn(1, 128, 128, 3)
-    flow = torch.zeros(1, 128, 128, 2) # Zero flow should be identity
-    flow[..., 0] = 10.0 # Shift 10px right
-
-    try:
-        warped = ofu.apply_flow(img, flow)
-        print(f"Warped shape: {warped.shape}")
-        if warped.shape == img.shape:
-            print("ofu.apply_flow SUCCESS!")
-        else:
-            print(f"ofu.apply_flow FAILED: Got {warped.shape}")
-    except Exception as e:
-        print(f"ofu.apply_flow FAILED with error: {e}")
-        import traceback
-        traceback.print_exc()
-
-if __name__ == "__main__":
-    test_permute_and_apply()
+    image = torch.rand(1, 32, 32, 3)
+    flow = torch.zeros(1, 32, 32, 2)
+    warped = ofu.apply_flow(image, flow)
+    assert warped.shape == image.shape
+    assert torch.allclose(warped, image, atol=1e-5)

tests/test_procedural.py

@@ -4,11 +4,12 @@ import os
 # Add the parent directory to sys.path so we can import more_math
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
+from more_math.Parser.antlr_router import get_antlr_modules
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
 from antlr4 import CommonTokenStream, InputStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 def evaluate(expression, variables=None):
     if variables is None:
         variables = {}
@@ -20,44 +21,28 @@ def evaluate(expression, variables=None):
     visitor = UnifiedMathVisitor(variables)
     return visitor.visit(tree)
 
-if __name__ == "__main__":
-    # 1. Test If-Else
+def test_if_else():
     expr_if = "a = 1; b = 0; if (a > 0) { b = 10; } else { b = 20; }; b"
-    res_if = evaluate(expr_if)
-    print(f"If-Else Test: b={res_if} (expected 10.0)")
+    assert evaluate(expr_if) == 10.0
 
-    # 2. Test While Loop
+
+def test_while_loop():
     expr_while = "i = 0; while (i < 5) { i = i + 1; }; i"
-    res_while = evaluate(expr_while)
-    print(f"While Loop Test: i={res_while} (expected 5.0)")
+    assert evaluate(expr_while) == 5.0
 
-    # 3. Test Function with Return and Recursion
-    # Fibonacci: f(n) -> { if(n <= 1) return n; return f(n-1) + f(n-2); }; f(7)
-    # Fib(7): 0, 1, 1, 2, 3, 5, 8, 13
+
+def test_recursive_function():
     expr_fib = "f(n) -> { if(n <= 1) return n; return f(n-1) + f(n-2); }; f(7)"
-    res_fib = evaluate(expr_fib)
-    print(f"Recursive Fib Test: f(7)={res_fib} (expected 13.0)")
+    assert evaluate(expr_fib) == 13.0
 
-    # 4. Test Isolated Block Scoping
-    # Variables defined inside blocks should not leak to outer scope
+
+def test_block_scope_removes_new_variables_but_preserves_assignments():
     expr_scope = "x = 1; { x = 2; y = 10; }; x"
-    # x should be modified to 2 inside block, but reverts to 1 after block
-    # y should not be accessible outside block
-    res_scope = evaluate(expr_scope)
-    print(f"Isolated Scope Test: x={res_scope} (expected 1.0)")
-    # 4b. Test that modifications to existing variables persist
+    variables = {}
+    assert evaluate(expr_scope, variables) == 2.0
+    assert "y" not in variables
+
+
+def test_existing_variable_assignment_persists_after_block():
     expr_scope2 = "x = 1; { x = x + 1; }; x"
-    res_scope2 = evaluate(expr_scope2)
-    print(f"Modified Variable Test: x={res_scope2} (expected 2.0)")
-
-    # 5. Test Return from top level
-    expr_top_return = "a = 10; return a + 1; 20"
-    res_top_return = evaluate(expr_top_return)
-    print(f"Top-level Return Test: {res_top_return} (expected 11.0)")
-
-    # 6. Test Deep Recursion with Procedural logic
-    # f(i) -> { if (i <= 0) return depth; return f(i-1); }; f(2000)
-    # Recursion limit is 10000, so 2000 should pass.
-    expr_deep = "f(i) -> { if (i <= 0) return depth; return f(i-1); }; f(2000)"
-    res_deep = evaluate(expr_deep)
-    print(f"Deep Recursion Test: f(2000)={res_deep} (expected 2001.0)")
+    assert evaluate(expr_scope2) == 2.0

tests/test_quantile_ops.py

@@ -10,10 +10,11 @@ if _project_root not in sys.path:
     sys.path.insert(0, _project_root)
 
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from antlr4 import InputStream, CommonTokenStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 def parse_and_visit(expr_str, variables):
     lexer = MathExprLexer(InputStream(expr_str))
     stream = CommonTokenStream(lexer)

tests/test_shorthands.py

@@ -10,10 +10,11 @@ if _project_root not in sys.path:
     sys.path.insert(0, _project_root)
 
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from antlr4 import InputStream, CommonTokenStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 def parse_and_visit(expr_str, variables):
     lexer = MathExprLexer(InputStream(expr_str))
     stream = CommonTokenStream(lexer)

tests/test_signal_stack.py

@@ -1,10 +1,11 @@
 import torch
 import sys
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from antlr4 import InputStream, CommonTokenStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 def parse_and_visit(expr, variables, state_storage=None):
     from antlr4.error.ErrorListener import ErrorListener
     class ThrowingErrorListener(ErrorListener):
@@ -26,8 +27,8 @@ def parse_and_visit(expr, variables, state_storage=None):
     visitor = UnifiedMathVisitor(variables, (1,), state_storage=state_storage)
     return visitor.visit(tree)
 
-def test_all():
-    with open("tests/test_debug.log", "w") as f_log:
+def test_all(tmp_path):
+    with (tmp_path / "test_debug.log").open("w") as f_log:
         def log(msg):
             print(msg)
             f_log.write(msg + "\n")
@@ -70,8 +71,8 @@ def test_all():
         assert parse_and_visit("cnt([1, 2, 3])", {}) == 3.0
         assert parse_and_visit("length([1, 2, 3])", {}) == 3.0
 
-        t = torch.randn(5, 10)
-        assert parse_and_visit("count(t)", {"t": t}) == 5.0
+        t = torch.arange(50).reshape(5, 10)
+        assert parse_and_visit("count(t)", {"t": t}) == 50
         assert parse_and_visit("count(42)", {}) == 1.0
 
         # 6. Tensor Function

tests/test_slicing.py

@@ -11,10 +11,11 @@ if _project_root not in sys.path:
     sys.path.insert(0, _project_root)
 
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from antlr4 import InputStream, CommonTokenStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 
 def parse_and_visit(expr_str, variables):
     """Helper to parse and visit an expression."""

tests/test_stats.py

@@ -10,10 +10,11 @@ if _project_root not in sys.path:
     sys.path.insert(0, _project_root)
 
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
 from antlr4 import InputStream, CommonTokenStream
 
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
+
 def parse_and_visit(expr_str, variables):
     lexer = MathExprLexer(InputStream(expr_str))
     stream = CommonTokenStream(lexer)

tests/test_unified_math.py

@@ -2,6 +2,7 @@ import sys
 import os
 import torch
 import math
+import pytest
 
 # Ensure we can import the module
 _here = os.path.abspath(os.path.dirname(__file__))
@@ -12,8 +13,10 @@ if _project_root not in sys.path:
 
 # Import visitor
 from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
-from more_math.Parser.MathExprLexer import MathExprLexer
-from more_math.Parser.MathExprParser import MathExprParser
+from more_math.Parser.antlr_router import get_antlr_modules
+from more_math.helper_functions import parse_expr
+
+MathExprLexer, MathExprParser, _ = get_antlr_modules()
 from antlr4 import InputStream, CommonTokenStream
 
 
@@ -34,6 +37,22 @@ def parse_and_visit(expr_str, variables):
     return visitor.visit(tree)
 
 
+def test_csharp_operator_precedence():
+    assert parse_and_visit("a + b << c", {"a": 1, "b": 2, "c": 1}) == 6
+    assert parse_and_visit("2 == 1 < 2", {}) == 0.0
+
+
+def test_rgb_conversion_dispatch():
+    assert parse_and_visit("int_to_rgb(16711680)", {}) == [255 / 256, 0, 0]
+    assert parse_and_visit("rgb_to_int([1, 0, 0])", {}) == 0xFF0000
+
+
+def test_map_rejects_more_than_three_coordinates():
+    parse_expr("map(t, x, y, z)")
+    with pytest.raises(ValueError, match="Syntax error"):
+        parse_expr("map(t, x, y, z, w)")
+
+
 def test_scalar_ops():
     vars = {"a": 2.0, "b": 3.0}