337 lines
15 KiB
Python
337 lines
15 KiB
Python
import torch
|
|
import sys
|
|
import os
|
|
import math
|
|
import pytest
|
|
|
|
# Ensure test runner usage (Visual Studio) can import the package regardless of working dir.
|
|
_here = os.path.abspath(os.path.dirname(__file__))
|
|
_project_root = os.path.abspath(os.path.join(_here, os.pardir))
|
|
_comfy_root = os.path.abspath(os.path.join(_project_root, os.pardir, os.pardir))
|
|
for _path in (_comfy_root, _project_root):
|
|
if _path not in sys.path:
|
|
sys.path.insert(0, _path)
|
|
|
|
from more_math.Parser.UnifiedMathVisitor import UnifiedMathVisitor
|
|
from more_math.Parser.inbuilt_symbols import INBUILT_FUNCTIONS
|
|
from more_math.helper_functions import parse_expr
|
|
|
|
def eval_expr(expr, variables={}, shape=(1,)):
|
|
tree = parse_expr(expr)
|
|
visitor = UnifiedMathVisitor(variables, shape)
|
|
return visitor.visit(tree)
|
|
|
|
def test_logic_functions():
|
|
assert eval_expr("any([0, 0, 1])") == 1.0
|
|
assert eval_expr("any([0, 0, 0])") == 0.0
|
|
assert eval_expr("all([1, 1, 1])") == 1.0
|
|
assert eval_expr("all([1, 0, 1])") == 0.0
|
|
|
|
t_any = torch.tensor([0.0, 1.0, 0.0])
|
|
assert eval_expr("any(T)", {"T": t_any}) == 1.0
|
|
|
|
def test_spatial_functions():
|
|
# Remap
|
|
r = eval_expr("remap(0.5, 0, 1, 10, 20)")
|
|
print(f"DEBUG: remap result: {r} type: {type(r)}")
|
|
assert abs(r - 15.0) < 1e-5
|
|
|
|
# Distance
|
|
d = eval_expr("dist(0,0, 3,4)")
|
|
print(f"DEBUG: dist result: {d} type: {type(d)}")
|
|
assert abs(d - 5.0) < 1e-5
|
|
|
|
# Edge
|
|
t_img = torch.zeros((1, 32, 32, 1))
|
|
t_img[:, 16:, :, :] = 1.0
|
|
res_edge = eval_expr("edge(I)", {"I": t_img})
|
|
print(f"DEBUG: edge shape: {res_edge.shape} max: {res_edge.max()}")
|
|
assert res_edge.shape == t_img.shape
|
|
assert res_edge.max() > 0
|
|
|
|
def test_stats_functions():
|
|
# Median
|
|
assert eval_expr("median([1, 10, 2, 5, 3])") == 3
|
|
|
|
# Mode
|
|
# Note: list mode uses Counter, returns first most common
|
|
assert eval_expr("mode([1, 2, 2, 3])") == 2
|
|
|
|
# Cumsum/Cumprod
|
|
t = torch.tensor([1, 2, 3], dtype=torch.float32)
|
|
assert torch.allclose(eval_expr("cumsum(T)", {"T": t}), torch.tensor([1, 3, 6], dtype=torch.float32))
|
|
assert torch.allclose(eval_expr("cumprod(T)", {"T": t}), torch.tensor([1, 2, 6], dtype=torch.float32))
|
|
|
|
def test_easing_functions():
|
|
# lerp(0, 10, 0.5) = 5
|
|
# cubic_ease(0, 10, 0.5) = 5 (inflection point)
|
|
assert abs(eval_expr("cubic_ease(0, 10, 0.5)") - 5.0) < 1e-5
|
|
# cubic_ease(0, 10, 0.1) should be less than lerp(0, 10, 0.1) = 1.0
|
|
assert eval_expr("cubic_ease(0, 10, 0.1)") < 1.0
|
|
|
|
# Sine ease
|
|
assert abs(eval_expr("sine_ease(0, 10, 0.5)") - 5.0) < 1e-5
|
|
|
|
# Smooterstep
|
|
assert abs(eval_expr("smootherstep(0.5, 0, 1)") - 0.5) < 1e-5
|
|
|
|
def test_lerp_fix():
|
|
# Test w as list
|
|
# lerp(0, 10, [0, 0.5, 1]) -> [0, 5, 10]
|
|
res = eval_expr("lerp(0, 10, [0, 0.5, 1])")
|
|
assert res == [0.0, 5.0, 10.0]
|
|
|
|
# Test a, b as lists matching w
|
|
res2 = eval_expr("lerp([0, 10, 20], [10, 20, 30], [0.5, 0.5, 0.5])")
|
|
# Expected: [5, 15, 25]
|
|
assert res2 == [5.0, 15.0, 25.0]
|
|
|
|
def test_added_math_and_logical_functions():
|
|
added_functions = {
|
|
"logical_not", "logical_and", "logical_or", "degrees", "radians", "log2", "hypot"
|
|
}
|
|
assert added_functions <= INBUILT_FUNCTIONS
|
|
|
|
assert eval_expr("logical_not(0)") == 1.0
|
|
assert eval_expr("logical_not(-2)") == 0.0
|
|
assert eval_expr("logical_and(2, 1)") == 1.0
|
|
assert eval_expr("logical_and(2, 0)") == 0.0
|
|
assert eval_expr("logical_or(0, -1)") == 1.0
|
|
assert eval_expr("logical_or(0, 0)") == 0.0
|
|
|
|
assert math.isclose(eval_expr("degrees(pi)"), 180.0, rel_tol=1e-6)
|
|
assert math.isclose(eval_expr("radians(180)"), math.pi, rel_tol=1e-6)
|
|
assert eval_expr("log2(8)") == 3.0
|
|
assert eval_expr("hypot(3, 4)") == 5.0
|
|
|
|
assert eval_expr("logical_and([1, 0], [1, 1])") == [1.0, 0.0]
|
|
assert eval_expr("logical_or([0, 0], [1, 0])") == [1.0, 0.0]
|
|
assert eval_expr("logical_not([0, 1])") == [1.0, 0.0]
|
|
assert torch.allclose(eval_expr("log2(T)", {"T": torch.tensor([1.0, 2.0, 8.0])}), torch.tensor([0.0, 1.0, 3.0]))
|
|
assert torch.allclose(eval_expr("hypot(A, B)", {"A": torch.tensor([3.0, 5.0]), "B": torch.tensor([4.0, 12.0])}), torch.tensor([5.0, 13.0]))
|
|
assert torch.allclose(eval_expr("logical_or(A, B)", {"A": torch.tensor([0.0, 1.0]), "B": torch.tensor([0.0, 0.0])}), torch.tensor([0.0, 1.0]))
|
|
assert torch.allclose(eval_expr("degrees(T)", {"T": torch.tensor([0.0, math.pi])}), torch.tensor([0.0, 180.0]))
|
|
|
|
def test_tensor_constructors_and_one_hot():
|
|
template = torch.empty((), dtype=torch.float64)
|
|
zeros = eval_expr("zeros([2, 3], T)", {"T": template})
|
|
assert zeros.shape == (2, 3)
|
|
assert zeros.dtype == torch.float64
|
|
assert torch.count_nonzero(zeros) == 0
|
|
assert torch.equal(eval_expr("ones([2])"), torch.ones(2))
|
|
assert torch.equal(eval_expr("eye(3)"), torch.eye(3))
|
|
assert torch.equal(eval_expr("eye(2, T)", {"T": template}), torch.eye(2, dtype=torch.float64))
|
|
assert torch.equal(
|
|
eval_expr("one_hot([0, 2], 3)"),
|
|
torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
|
|
)
|
|
assert "one_hot" in INBUILT_FUNCTIONS
|
|
with pytest.raises(ValueError, match="outside the class range"):
|
|
eval_expr("one_hot([3], 3)")
|
|
with pytest.raises(ValueError, match="must be integers"):
|
|
eval_expr("one_hot([0.5], 3)")
|
|
|
|
def test_loss_metrics_and_psnr():
|
|
a = torch.tensor([1.0, 3.0])
|
|
b = torch.tensor([2.0, 1.0])
|
|
variables = {"a": a, "b": b}
|
|
assert math.isclose(eval_expr("mse(a, b)", variables).item(), 2.5)
|
|
assert math.isclose(eval_expr("mse(a, b, 'sum')", variables).item(), 5.0)
|
|
assert torch.equal(eval_expr("mse(a, b, 'none')", variables), torch.tensor([1.0, 4.0]))
|
|
assert math.isclose(eval_expr("mae(a, b)", variables).item(), 1.5)
|
|
assert math.isclose(eval_expr("mae(a, b, 'sum')", variables).item(), 3.0)
|
|
assert math.isclose(eval_expr("psnr(a, b)", variables).item(), 10.0 * math.log10(0.4), rel_tol=1e-6)
|
|
assert math.isinf(eval_expr("psnr(a, a)", variables).item())
|
|
assert "mse" in INBUILT_FUNCTIONS and "mae" in INBUILT_FUNCTIONS and "psnr" in INBUILT_FUNCTIONS
|
|
with pytest.raises(ValueError, match="reduction must be"):
|
|
eval_expr("mse(a, b, 'median')", variables)
|
|
with pytest.raises(ValueError, match="data_range must be positive"):
|
|
eval_expr("psnr(a, b, 0)", variables)
|
|
|
|
def test_leaky_relu_and_softsign():
|
|
assert eval_expr("leaky_relu([-2, 0, 2])") == [-0.02, 0.0, 2.0]
|
|
assert eval_expr("leaky_relu(-2, 0.2)") == -0.4
|
|
assert eval_expr("softsign([-2, 0, 2])") == [-2.0 / 3.0, 0.0, 2.0 / 3.0]
|
|
assert "leaky_relu" in INBUILT_FUNCTIONS and "softsign" in INBUILT_FUNCTIONS
|
|
|
|
def test_tnorm_normalizes_lists():
|
|
torch.testing.assert_close(torch.tensor(eval_expr("tnorm([3, 4])")), torch.tensor([0.6, 0.8]))
|
|
torch.testing.assert_close(torch.tensor(eval_expr("tnorm([[3, 4], [0, 0]])")), torch.tensor([[0.6, 0.8], [0.0, 0.0]]))
|
|
torch.testing.assert_close(torch.tensor(eval_expr("tnorm([[3, 0], [0, 4]], 0)")), torch.tensor([[1.0, 0.0], [0.0, 1.0]]))
|
|
|
|
def test_silu_activation():
|
|
expected = torch.nn.functional.silu(torch.tensor([-2.0, 0.0, 2.0])).tolist()
|
|
assert eval_expr("silu([-2, 0, 2])") == pytest.approx(expected)
|
|
assert eval_expr("silu(1000)") == pytest.approx(1000.0)
|
|
|
|
values = torch.tensor([-3.0, 0.0, 2.0])
|
|
torch.testing.assert_close(eval_expr("silu(x)", {"x": values}), torch.nn.functional.silu(values))
|
|
assert "silu" in INBUILT_FUNCTIONS
|
|
|
|
def test_group_norm_matches_pytorch():
|
|
values = torch.arange(24, dtype=torch.float32).reshape(2, 4, 3)
|
|
weight = torch.tensor([1.0, 0.5, 2.0, 3.0])
|
|
bias = torch.tensor([0.0, 1.0, -1.0, 2.0])
|
|
variables = {"x": values, "weight": weight, "bias": bias}
|
|
|
|
torch.testing.assert_close(
|
|
eval_expr("group_norm(x, 2)", variables),
|
|
torch.nn.functional.group_norm(values, 2),
|
|
)
|
|
torch.testing.assert_close(
|
|
eval_expr("group_norm(x, 2, weight, bias, 0.01)", variables),
|
|
torch.nn.functional.group_norm(values, 2, weight, bias, eps=0.01),
|
|
)
|
|
assert "group_norm" in INBUILT_FUNCTIONS
|
|
|
|
with pytest.raises(ValueError, match="divide the channel count"):
|
|
eval_expr("group_norm(x, 3)", variables)
|
|
with pytest.raises(ValueError, match="weight must have shape"):
|
|
eval_expr("group_norm(x, 2, [1, 2])", variables)
|
|
with pytest.raises(ValueError, match="eps must be a positive finite scalar"):
|
|
eval_expr("group_norm(x, 2, none, none, 0)", variables)
|
|
|
|
def test_group_norm_preserves_float8_dtype():
|
|
float8_dtypes = [
|
|
getattr(torch, name, None)
|
|
for name in ("float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
|
|
]
|
|
for dtype in (dtype for dtype in float8_dtypes if dtype is not None):
|
|
values = torch.randn(2, 4, 3).to(dtype)
|
|
result = eval_expr("group_norm(x, 2)", {"x": values})
|
|
assert result.dtype == dtype
|
|
|
|
def test_slerp_interpolates_direction_and_magnitude():
|
|
a = torch.tensor([2.0, 0.0])
|
|
b = torch.tensor([0.0, 4.0])
|
|
result = eval_expr("slerp(a, b, 0.5)", {"a": a, "b": b})
|
|
torch.testing.assert_close(result, torch.tensor([3.0 / math.sqrt(2.0), 3.0 / math.sqrt(2.0)]))
|
|
torch.testing.assert_close(eval_expr("slerp(a, b, 0)", {"a": a, "b": b}), a)
|
|
torch.testing.assert_close(eval_expr("slerp(a, b, 1)", {"a": a, "b": b}), b)
|
|
|
|
batch_a = torch.tensor([[1.0, 0.0], [2.0, 0.0]])
|
|
batch_b = torch.tensor([[0.0, 1.0], [0.0, 2.0]])
|
|
torch.testing.assert_close(
|
|
eval_expr("slerp(a, b, t, 1)", {"a": batch_a, "b": batch_b, "t": torch.tensor([0.0, 1.0])}),
|
|
torch.tensor([[1.0, 0.0], [0.0, 2.0]]),
|
|
)
|
|
|
|
whole_a = torch.tensor([[1.0, 0.0], [0.0, 0.0]])
|
|
whole_b = torch.tensor([[0.0, 0.0], [0.0, 1.0]])
|
|
torch.testing.assert_close(
|
|
eval_expr("slerp(a, b, 0.5)", {"a": whole_a, "b": whole_b}),
|
|
torch.tensor([[1.0 / math.sqrt(2.0), 0.0], [0.0, 1.0 / math.sqrt(2.0)]]),
|
|
)
|
|
|
|
a_by_column = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
|
|
b_by_column = torch.tensor([[0.0, 2.0], [2.0, 0.0]])
|
|
torch.testing.assert_close(
|
|
eval_expr("slerp(a, b, 0.5, 0)", {"a": a_by_column, "b": b_by_column}),
|
|
torch.full((2, 2), 1.5 / math.sqrt(2.0)),
|
|
)
|
|
|
|
antipodal = eval_expr("slerp([1, 0], [-4, 0], 0.5)")
|
|
assert torch.isfinite(antipodal).all()
|
|
torch.testing.assert_close(torch.linalg.vector_norm(antipodal), torch.tensor(2.5))
|
|
torch.testing.assert_close(eval_expr("slerp([0, 0], [0, 4], 0.5)"), torch.tensor([0.0, 2.0]))
|
|
assert "slerp" in INBUILT_FUNCTIONS
|
|
|
|
def test_slerp_validates_broadcasting_and_preserves_float8_dtype():
|
|
with pytest.raises(ValueError, match="inputs must be broadcastable"):
|
|
eval_expr("slerp([1, 2], [1, 2, 3], 0.5)")
|
|
with pytest.raises(ValueError, match="t must broadcast"):
|
|
eval_expr("slerp([[1, 0], [0, 1]], [[0, 1], [1, 0]], t, 1)", {"t": torch.tensor([1.0, 2.0, 3.0])})
|
|
|
|
float8_dtypes = [
|
|
getattr(torch, name, None)
|
|
for name in ("float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
|
|
]
|
|
for dtype in (dtype for dtype in float8_dtypes if dtype is not None):
|
|
values = torch.randn(2, 3).to(dtype)
|
|
result = eval_expr("slerp(a, b, 0.5)", {"a": values, "b": values})
|
|
assert result.dtype == dtype
|
|
|
|
def test_slerp_and_lerp_with_list_of_times_concat_to_batch():
|
|
a = torch.randn(1, 3, 4)
|
|
b = torch.randn(1, 3, 4)
|
|
c = torch.randn(1, 3, 4)
|
|
times = [0.0, 0.25, 1.0]
|
|
variables = {"a": a, "b": b, "c": c}
|
|
|
|
result = eval_expr("concat(slerp(a, b, [0, 0.25, 1]), 0)", variables)
|
|
expected = torch.cat([eval_expr(f"slerp(a, b, {t})", variables) for t in times])
|
|
torch.testing.assert_close(result, expected)
|
|
torch.testing.assert_close(result[0:1], a)
|
|
torch.testing.assert_close(result[2:3], b)
|
|
|
|
# nested calls pair each time with its own result
|
|
result = eval_expr("w = [0, 0.25, 1]; concat(slerp(slerp(a, b, w), slerp(b, c, w), w), 0)", variables)
|
|
expected = torch.cat([eval_expr(f"slerp(slerp(a, b, {t}), slerp(b, c, {t}), {t})", variables) for t in times])
|
|
torch.testing.assert_close(result, expected)
|
|
|
|
result = eval_expr("w = [0, 0.25, 1]; concat(lerp(lerp(a, b, w), lerp(b, c, w), w), 0)", variables)
|
|
expected = torch.cat([torch.lerp(torch.lerp(a, b, t), torch.lerp(b, c, t), t) for t in times])
|
|
torch.testing.assert_close(result, expected)
|
|
|
|
with pytest.raises(ValueError, match="2 values for 3 times"):
|
|
eval_expr("slerp(slerp(a, b, [0, 1]), b, [0, 0.5, 1])", variables)
|
|
|
|
def test_eases_and_interp_with_list_of_times():
|
|
a = torch.randn(1, 3, 4)
|
|
b = torch.randn(1, 3, 4)
|
|
variables = {"a": a, "b": b}
|
|
times = [0.0, 0.3, 0.8, 1.0]
|
|
for name in ("cubic_ease", "sine_ease", "elastic_ease"):
|
|
result = eval_expr(f"concat({name}(a, b, [0, 0.3, 0.8, 1]), 0)", variables)
|
|
expected = torch.cat([eval_expr(f"{name}(a, b, {t})", variables) for t in times])
|
|
torch.testing.assert_close(result, expected)
|
|
assert eval_expr("sine_ease(0, 2, [0, 0.5, 1])") == pytest.approx([0.0, 1.0, 2.0])
|
|
|
|
keys = {"a": a, "b": b, "c": torch.randn(1, 3, 4)}
|
|
for mode in ("linear", "smooth", "step"):
|
|
result = eval_expr(f'concat(interp([0, 0.25, 0.5, 0.9], [0, 0.5, 1], [a, b, c], "{mode}"), 0)', keys)
|
|
expected = torch.cat([eval_expr(f'interp({t}, [0, 0.5, 1], [a, b, c], "{mode}")', keys) for t in (0, 0.25, 0.5, 0.9)])
|
|
torch.testing.assert_close(result, expected)
|
|
assert eval_expr("interp([0, 12, 24], [0, 24], [1, 3])") == [1.0, 2.0, 3.0]
|
|
|
|
def test_concat_single_list_argument():
|
|
a = torch.ones(1, 2)
|
|
b = torch.zeros(2, 2)
|
|
torch.testing.assert_close(eval_expr("concat([a, b], 0)", {"a": a, "b": b}), torch.cat([a, b]))
|
|
torch.testing.assert_close(eval_expr("concat([a, a], 1)", {"a": a}), torch.ones(1, 4))
|
|
assert eval_expr("concat([1, 2], [3], 0)") == [1, 2, 3]
|
|
|
|
def test_sinc_and_erfc():
|
|
assert {"sinc", "erfc"} <= INBUILT_FUNCTIONS
|
|
assert eval_expr("sinc(0)") == 1.0
|
|
assert abs(eval_expr("sinc(0.5)") - 2.0 / math.pi) < 1e-6
|
|
assert abs(eval_expr("sinc(1)")) < 1e-6
|
|
assert eval_expr("erfc(0)") == 1.0
|
|
assert math.isclose(eval_expr("erfc(1)"), math.erfc(1), rel_tol=1e-6)
|
|
torch.testing.assert_close(
|
|
eval_expr("sinc(x)", {"x": torch.tensor([0.0, 0.5, 1.0])}),
|
|
torch.sinc(torch.tensor([0.0, 0.5, 1.0])),
|
|
)
|
|
torch.testing.assert_close(
|
|
eval_expr("erfc(x)", {"x": torch.tensor([-1.0, 0.0, 1.0])}),
|
|
torch.special.erfc(torch.tensor([-1.0, 0.0, 1.0])),
|
|
)
|
|
assert eval_expr("sinc([0, 0.5])") == [1.0, 2.0 / math.pi]
|
|
float8_dtypes = [
|
|
getattr(torch, name, None)
|
|
for name in ("float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
|
|
]
|
|
for dtype in (dtype for dtype in float8_dtypes if dtype is not None):
|
|
values = torch.tensor([0.0, 0.5, 1.0]).to(dtype)
|
|
assert eval_expr("sinc(x)", {"x": values}).dtype == dtype
|
|
assert eval_expr("erfc(x)", {"x": values}).dtype == dtype
|
|
|
|
if __name__ == "__main__":
|
|
test_logic_functions()
|
|
test_spatial_functions()
|
|
test_stats_functions()
|
|
test_easing_functions()
|
|
test_lerp_fix()
|
|
test_sinc_and_erfc()
|
|
print("All new function tests passed!")
|