161 lines
7.1 KiB
Python
161 lines
7.1 KiB
Python
import torch
|
|
import sys
|
|
import os
|
|
import math
|
|
import pytest
|
|
|
|
# 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.antlr_router import get_antlr_modules
|
|
from more_math.Parser.inbuilt_symbols import INBUILT_FUNCTIONS, INBUILT_KEYWORDS
|
|
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)
|
|
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
|
|
|
|
@pytest.mark.parametrize("function", ("mean", "std", "var"))
|
|
@pytest.mark.parametrize("dimension", ("0", "-1", "[0, 1]"))
|
|
def test_statistic_uses_requested_dimensions(function, dimension):
|
|
values = torch.tensor([[1.0, 3.0, 5.0], [2.0, 4.0, 8.0]])
|
|
if dimension == "[0, 1]":
|
|
expected = getattr(torch, function)(values.float(), dim=(0, 1))
|
|
else:
|
|
expected = getattr(torch, function)(values.float(), dim=int(dimension))
|
|
|
|
result = parse_and_visit(f"{function}(t, {dimension})", {"t": values})
|
|
|
|
if isinstance(expected, torch.Tensor) and expected.numel() == 1:
|
|
expected = expected.item()
|
|
if isinstance(result, torch.Tensor):
|
|
torch.testing.assert_close(result, expected)
|
|
else:
|
|
assert math.isclose(result, expected, rel_tol=1e-5)
|
|
|
|
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
|
|
|
|
def test_isclose_and_allclose_are_functions():
|
|
a = torch.tensor([1.0, 2.0])
|
|
b = torch.tensor([1.0, 2.000001])
|
|
|
|
assert "isclose" in INBUILT_FUNCTIONS
|
|
assert "allclose" in INBUILT_FUNCTIONS
|
|
assert "isclose" not in INBUILT_KEYWORDS
|
|
assert "allclose" not in INBUILT_KEYWORDS
|
|
torch.testing.assert_close(parse_and_visit("isclose(a, b)", {"a": a, "b": b}), torch.tensor([1.0, 1.0]))
|
|
assert parse_and_visit("allclose(a, b)", {"a": a, "b": b}) == 1.0
|
|
|
|
def test_order_statistics_support_dimensions():
|
|
values = torch.tensor([[1.0, 3.0, 3.0], [2.0, 2.0, 8.0]])
|
|
variables = {"t": values}
|
|
|
|
torch.testing.assert_close(parse_and_visit("median(t, 0)", variables), torch.tensor([1.0, 2.0, 3.0]))
|
|
torch.testing.assert_close(parse_and_visit("mode(t, 1)", variables), torch.tensor([3.0, 2.0]))
|
|
torch.testing.assert_close(parse_and_visit("quantile(t, 0.5, 0)", variables), torch.tensor([1.5, 2.5, 5.5]))
|
|
torch.testing.assert_close(parse_and_visit("percentile(t, 50, 1)", variables), torch.tensor([3.0, 2.0]))
|
|
torch.testing.assert_close(parse_and_visit("quartile(t, 2, 1)", variables), torch.tensor([3.0, 2.0]))
|
|
|
|
torch.testing.assert_close(parse_and_visit("argmin(t, 0, 1)", variables), torch.tensor([0, 0]))
|
|
torch.testing.assert_close(parse_and_visit("argmax(t, 0, 1)", variables), torch.tensor([1, 2]))
|
|
torch.testing.assert_close(parse_and_visit("topk_ind(t, 1, 1)", variables), torch.tensor([[1], [2]]))
|
|
torch.testing.assert_close(parse_and_visit("botk_ind(t, 1, 1)", variables), torch.tensor([[0], [0]]))
|
|
torch.testing.assert_close(
|
|
parse_and_visit("topk(t, 1, 1)", variables),
|
|
torch.tensor([[0.0, 3.0, 0.0], [0.0, 0.0, 8.0]]),
|
|
)
|
|
|
|
def test_weighted_mean_and_losses():
|
|
values = torch.tensor([[1.0, 3.0], [2.0, 4.0]])
|
|
weights = torch.tensor([[1.0, 3.0], [3.0, 1.0]])
|
|
variables = {"x": values, "w": weights}
|
|
torch.testing.assert_close(parse_and_visit("weighted_mean(x, w)", variables), torch.tensor(2.5))
|
|
torch.testing.assert_close(parse_and_visit("weighted_mean(x, w, 1)", variables), torch.tensor([2.5, 2.5]))
|
|
torch.testing.assert_close(parse_and_visit("weighted_mean(x, w, 0)", variables), torch.tensor([1.75, 3.25]))
|
|
|
|
prediction = torch.tensor([0.0, 2.0])
|
|
target = torch.tensor([1.0, 0.0])
|
|
torch.testing.assert_close(
|
|
parse_and_visit("huber_loss(prediction, target)", {"prediction": prediction, "target": target}),
|
|
torch.tensor(1.0),
|
|
)
|
|
assert parse_and_visit("huber_loss(0, 2, 0.5, 'sum')", {}) == 0.875
|
|
with pytest.raises(ValueError, match="delta must be a positive"):
|
|
parse_and_visit("huber_loss(0, 1, 0)", {})
|
|
with pytest.raises(ValueError, match="weights must have a non-zero sum"):
|
|
parse_and_visit("weighted_mean([1, 2], [1, -1])", {})
|
|
|
|
logits = torch.tensor([[2.0, 0.0], [0.0, 2.0]])
|
|
labels = torch.tensor([0, 1])
|
|
expected = torch.nn.functional.cross_entropy(logits, labels)
|
|
torch.testing.assert_close(
|
|
parse_and_visit("cross_entropy(logits, labels)", {"logits": logits, "labels": labels}),
|
|
expected,
|
|
)
|
|
probabilities = torch.nn.functional.one_hot(labels, num_classes=2).float()
|
|
torch.testing.assert_close(
|
|
parse_and_visit("cross_entropy(logits, probabilities, 1, 'sum')", {"logits": logits, "probabilities": probabilities}),
|
|
torch.nn.functional.cross_entropy(logits, probabilities, reduction="sum"),
|
|
)
|
|
|
|
logits_by_class = torch.tensor([[2.0, 0.0], [0.0, 2.0], [0.0, 0.0]])
|
|
torch.testing.assert_close(
|
|
parse_and_visit("cross_entropy(logits, labels, 0)", {"logits": logits_by_class, "labels": labels}),
|
|
torch.nn.functional.cross_entropy(logits_by_class.movedim(0, 1), labels),
|
|
)
|
|
|