From d04d386a1674acade85ca2c08125c2a8ee5db60e Mon Sep 17 00:00:00 2001 From: mcDandy Date: Mon, 12 Jan 2026 15:02:24 +0100 Subject: [PATCH] (ai) add test for moment --- tests/test_quantile_ops.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_quantile_ops.py b/tests/test_quantile_ops.py index a6fadac..1b8d7e8 100644 --- a/tests/test_quantile_ops.py +++ b/tests/test_quantile_ops.py @@ -118,6 +118,20 @@ def test_list_query(): assert res[1] == 20.0 assert res[2] == 40.0 +def test_moment(): + # moment(x, a, k) = 1/n * sum((x-a)^k) + t = torch.tensor([1.0, 2.0, 3.0]) + vars = {"t": t, "l": [1, 2, 3]} + + # First moment about 0 (mean) + assert parse_and_visit("moment(t, 0, 1)", vars) == 2.0 + assert parse_and_visit("moment(l, 0, 1)", vars) == 2.0 + + # Second moment about mean (variance) + # (1-2)^2 + (2-2)^2 + (3-2)^2 = 1 + 0 + 1 = 2. 2/3 = 0.666... + res = parse_and_visit("moment(t, 2, 2)", vars) + assert math.isclose(res, 2.0/3.0, rel_tol=1e-5) + if __name__ == "__main__": # If run as script try: @@ -129,6 +143,7 @@ if __name__ == "__main__": test_nd_tensor_query() test_large_tensor_fallback() test_list_query() + test_moment() print("All quantile_ops tests passed!") except Exception: import traceback