71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
import torch
|
|
import sys
|
|
import os
|
|
import math
|
|
|
|
# Ensure we can import the module
|
|
_here = os.path.abspath(os.path.dirname(__file__))
|
|
_project_root = os.path.abspath(os.path.join(_here, os.pardir))
|
|
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 antlr4 import InputStream, CommonTokenStream
|
|
|
|
def parse_and_visit(expr_str, variables):
|
|
lexer = MathExprLexer(InputStream(expr_str))
|
|
stream = CommonTokenStream(lexer)
|
|
parser = MathExprParser(stream)
|
|
tree = parser.expr()
|
|
shape = (1, 1, 1, 1) # Dummy shape
|
|
visitor = UnifiedMathVisitor(variables, shape)
|
|
return visitor.visit(tree)
|
|
|
|
def test_sum():
|
|
vars = {"t": torch.tensor([1.0, 2.0, 3.0]), "l": [1.0, 2.0, 3.0], "s": 5.0}
|
|
assert parse_and_visit("sum(t)", vars) == 6.0
|
|
assert parse_and_visit("sum(l)", vars) == 6.0
|
|
assert parse_and_visit("sum(s)", vars) == 5.0
|
|
|
|
def test_mean():
|
|
vars = {"t": torch.tensor([1.0, 2.0, 3.0]), "l": [1.0, 2.0, 3.0], "s": 5.0}
|
|
assert parse_and_visit("mean(t)", vars) == 2.0
|
|
assert parse_and_visit("mean(l)", vars) == 2.0
|
|
assert parse_and_visit("mean(s)", vars) == 5.0
|
|
|
|
def test_std():
|
|
t = torch.tensor([1.0, 5.0]) # mean=3, var=((2^2 + 2^2)/1) = 8. std = sqrt(8)=2.828...
|
|
vars = {"t": t, "l": [1.0, 5.0]}
|
|
|
|
res_t = parse_and_visit("std(t)", vars)
|
|
# Result might be float or tensor with .item()
|
|
val_t = res_t.item() if hasattr(res_t, 'item') else res_t
|
|
assert math.isclose(val_t, math.sqrt(8), rel_tol=1e-5)
|
|
assert math.isclose(parse_and_visit("std(l)", vars), math.sqrt(8), rel_tol=1e-5)
|
|
|
|
def test_var():
|
|
t = torch.tensor([1.0, 5.0])
|
|
vars = {"t": t, "l": [1.0, 5.0]}
|
|
res_t = parse_and_visit("var(t)", vars)
|
|
# Result might be float or tensor with .item()
|
|
val_t = res_t.item() if hasattr(res_t, 'item') else res_t
|
|
assert math.isclose(val_t, 8.0, rel_tol=1e-5)
|
|
assert math.isclose(parse_and_visit("var(l)", vars), 8.0, rel_tol=1e-5)
|
|
|
|
def test_dot():
|
|
a = torch.tensor([1.0, 2.0])
|
|
b = torch.tensor([3.0, 4.0])
|
|
vars = {"a": a, "b": b}
|
|
# dot = 1*3 + 2*4 = 3 + 8 = 11
|
|
assert parse_and_visit("dot(a, b)", vars) == 11.0
|
|
|
|
# Check flattening behavior (2D tensor)
|
|
a2 = torch.tensor([[1.0, 2.0]])
|
|
b2 = torch.tensor([[3.0], [4.0]]) # 2x1
|
|
vars = {"a2": a2, "b2": b2}
|
|
# a2 flat = [1, 2], b2 flat = [3, 4] -> dot=11
|
|
assert parse_and_visit("dot(a2, b2)", vars) == 11.0
|
|
|