From a925909e10957ba231b9dbd0be6a1d78fc3693f2 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Tue, 13 Oct 2026 15:58:07 +0200 Subject: [PATCH] AI: sstylistic and fixes of visitor --- more_math/Parser/UnifiedMathVisitor.py | 156 +++++++++++++--------- more_math/modelLikeCommon.py | 7 +- tests/test_unified_visitor_regressions.py | 90 +++++++++++++ 3 files changed, 187 insertions(+), 66 deletions(-) create mode 100644 tests/test_unified_visitor_regressions.py diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index 51935b4..bed61d3 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/more_math/Parser/UnifiedMathVisitor.py @@ -1,6 +1,5 @@ import time import os -from PIL.XVThumbImagePlugin import r import torch import math import inspect @@ -52,22 +51,49 @@ class ChangeTrackingDict(dict): self._changed_keys = set() self._tracking_enabled = True + @staticmethod + def _values_equal(left, right): + if isinstance(left, torch.Tensor) or isinstance(right, torch.Tensor): + if not isinstance(left, torch.Tensor) or not isinstance(right, torch.Tensor): + return False + try: + return torch.equal(left, right) + except RuntimeError: + return False + if isinstance(left, (list, tuple)) or isinstance(right, (list, tuple)): + return ( + isinstance(left, (list, tuple)) + and isinstance(right, (list, tuple)) + and len(left) == len(right) + and all(ChangeTrackingDict._values_equal(a, b) for a, b in zip(left, right)) + ) + try: + return bool(left == right) + except (TypeError, ValueError, RuntimeError): + return left is right + + @classmethod + def from_result(cls, values, source): + result = cls(values) + result._initial_keys = source._initial_keys.copy() + result._initial_values = source._initial_values.copy() + result._tracking_enabled = source._tracking_enabled + if result._tracking_enabled: + result._changed_keys = { + key for key, value in result.items() + if key not in result._initial_keys + or not cls._values_equal(value, result._initial_values[key]) + } + return result + def __setitem__(self, key, value): if self._tracking_enabled: - # Check if this is a modification (key exists and value changed) if key in self._initial_keys: - # Key existed before - check if value changed - if key in self._initial_values: - initial = self._initial_values[key] - # For tensors, compare identity and value - if isinstance(value, torch.Tensor) and isinstance(initial, torch.Tensor): - if value is not initial and not torch.equal(value, initial): - self._changed_keys.add(key) - else: - if value is not initial: - self._changed_keys.add(key) + if key in self._initial_values and self._values_equal(value, self._initial_values[key]): + self._changed_keys.discard(key) + else: + self._changed_keys.add(key) else: - # New key - mark as changed (can be added key) self._changed_keys.add(key) super().__setitem__(key, value) @@ -190,10 +216,10 @@ class UnifiedMathVisitor(MathExprVisitor): if any(set(value) != key_set for value in dictionaries[1:]): raise ValueError("Dictionary operations require matching keys") - # Determine result type: if any input is ChangeTrackingDict, create ChangeTrackingDict result_dict = {key: op(*(value[key] if self._is_math_dict(value) else value for value in values)) for key in keys} - if any(isinstance(value, ChangeTrackingDict) for value in dictionaries): - return ChangeTrackingDict(result_dict) + tracking_source = next((value for value in dictionaries if isinstance(value, ChangeTrackingDict)), None) + if tracking_source is not None: + return ChangeTrackingDict.from_result(result_dict, tracking_source) return MathDict(result_dict) def _where_value(self, condition, true_value, false_value): @@ -218,6 +244,45 @@ class UnifiedMathVisitor(MathExprVisitor): return torch.full(t,val,device=self.device) return torch.tensor(val, device=self.device) + @staticmethod + def _float8_dtypes(): + return { + dtype for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e4m3fnuz", None), + getattr(torch, "float8_e5m2", None), + getattr(torch, "float8_e5m2fnuz", None), + ) + if dtype is not None + } + + def _apply_binary_operation(self, a, b, torch_op, scalar_op): + float8_dtypes = self._float8_dtypes() + original_dtype = None + if self._is_plain_tensor(a) and a.dtype in float8_dtypes: + original_dtype = a.dtype + a = a.to(torch.float16) + if self._is_plain_tensor(b) and b.dtype in float8_dtypes: + if original_dtype is None: + original_dtype = b.dtype + b = b.to(torch.float16) + + result = torch_op(a, b).contiguous() if torch_op else scalar_op(a, b) + if original_dtype is not None and self._is_plain_tensor(result): + result = result.to(original_dtype) + return result + + def _apply_unary_operation(self, value, torch_op, scalar_op): + original_dtype = None + if self._is_plain_tensor(value) and value.dtype in self._float8_dtypes(): + original_dtype = value.dtype + value = value.to(torch.float16) + + result = torch_op(value).contiguous() if torch_op else scalar_op(value) + if original_dtype is not None and self._is_plain_tensor(result): + result = result.to(original_dtype) + return result + def _bin_op(self, a, b, torch_op, scalar_op, ctx): """ Generic binary operation handler. @@ -228,18 +293,18 @@ class UnifiedMathVisitor(MathExprVisitor): if self._is_math_dict(a) and not self._is_math_dict(b): if not self._is_tensor(b) and not self._is_list(b): # a is dict, b is scalar - apply op directly to all dict values - result = {key: (torch_op(a[key], b) if torch_op else scalar_op(a[key], b)) + result = {key: self._apply_binary_operation(a[key], b, torch_op, scalar_op) for key in a} if isinstance(a, ChangeTrackingDict): - return ChangeTrackingDict(result) + return ChangeTrackingDict.from_result(result, a) return MathDict(result) elif self._is_math_dict(b) and not self._is_math_dict(a): if not self._is_tensor(a) and not self._is_list(a): # b is dict, a is scalar - apply op directly to all dict values - result = {key: (torch_op(a, b[key]) if torch_op else scalar_op(a, b[key])) + result = {key: self._apply_binary_operation(a, b[key], torch_op, scalar_op) for key in b} if isinstance(b, ChangeTrackingDict): - return ChangeTrackingDict(result) + return ChangeTrackingDict.from_result(result, b) return MathDict(result) if self._is_math_dict(a) or self._is_math_dict(b): @@ -303,34 +368,7 @@ class UnifiedMathVisitor(MathExprVisitor): return NestedTensor([self._bin_op(a_comps, x, torch_op, scalar_op, ctx) for x in b_comps]) if self._is_tensor(a) or self._is_tensor(b): - orig_dtype = None - - # Promote float8 to float16 if needed - if self._is_tensor(a): - if hasattr(torch, "float8_e4m3fn") and a.dtype == torch.float8_e4m3fn: - orig_dtype = a.dtype - a = a.to(torch.float16) - elif hasattr(torch, "float8_e5m2") and a.dtype == torch.float8_e5m2: - orig_dtype = a.dtype - a = a.to(torch.float16) - - if self._is_tensor(b): - if hasattr(torch, "float8_e4m3fn") and b.dtype == torch.float8_e4m3fn: - orig_dtype = b.dtype if orig_dtype is None else orig_dtype - b = b.to(torch.float16) - elif hasattr(torch, "float8_e5m2") and b.dtype == torch.float8_e5m2: - orig_dtype = b.dtype if orig_dtype is None else orig_dtype - b = b.to(torch.float16) - - if torch_op: - res = torch_op(a, b).contiguous() - else: - res = scalar_op(a, b) - - if orig_dtype is not None and self._is_tensor(res): - res = res.to(orig_dtype) - - return res + return self._apply_binary_operation(a, b, torch_op, scalar_op) return scalar_op(a, b) except (ArithmeticError) as e: @@ -349,19 +387,7 @@ class UnifiedMathVisitor(MathExprVisitor): if self._is_list(a): return [self._unary_op(x, torch_op, scalar_op) for x in a] if self._is_tensor(a): - orig_dtype = None - if hasattr(torch, "float8_e4m3fn") and a.dtype == torch.float8_e4m3fn: - orig_dtype = a.dtype - a = a.to(torch.float16) - elif hasattr(torch, "float8_e5m2") and a.dtype == torch.float8_e5m2: - orig_dtype = a.dtype - a = a.to(torch.float16) - - res = torch_op(a).contiguous() if torch_op else scalar_op(a) - - if orig_dtype is not None and self._is_tensor(res): - res = res.to(orig_dtype) - return res + return self._apply_unary_operation(a, torch_op, scalar_op) return scalar_op(a) @@ -2654,10 +2680,15 @@ class UnifiedMathVisitor(MathExprVisitor): for e in ctx.exprList().expr(): args.append((yield e)) + if len(args) != len(lam.params): + raise ValueError( + f"{ctx.start.line}:{ctx.start.column}: Lambda expects {len(lam.params)} arguments, got {len(args)}" + ) + # připrav nový scope na základě uzávěrky new_vars = lam.env.copy() for i, p in enumerate(lam.params): - new_vars[p] = args[i] if i < len(args) else None + new_vars[p] = args[i] # push/pop scope stejným stylem jako pro pojmenované funkce self._scope_stack.append(self.variables) @@ -2747,7 +2778,6 @@ class UnifiedMathVisitor(MathExprVisitor): text = _to_str((yield ctx.expr(0))) font_name = _to_str((yield ctx.expr(1))) size = max(1, self._to_int((yield ctx.expr(2)), ctx,"text_image", strict=True)) - size = max(1, self._to_int((yield ctx.expr(2)), ctx,"text_image", strict=True)) max_width = self._to_int((yield ctx.expr(3)), ctx,"text_image", strict=True) if len(ctx.expr()) > 3 else 0 weight_val = self._to_int((yield ctx.expr(4)), ctx,"text_image", strict=True) if len(ctx.expr()) > 4 else 400 angle = _to_float((yield ctx.expr(5))) if len(ctx.expr()) > 5 else 0.0 diff --git a/more_math/modelLikeCommon.py b/more_math/modelLikeCommon.py index 9b39d1e..33f14ff 100644 --- a/more_math/modelLikeCommon.py +++ b/more_math/modelLikeCommon.py @@ -467,9 +467,10 @@ def calculate_patches_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=[], u return {} def weight_dict(name): - """Create a ChangeTrackingDict that monitors changes to weights""" - tracking_dict = ChangeTrackingDict({key: weights[name].get(key, torch.zeros_like(references[key])) for key in all_keys}) - return tracking_dict + values = {key: weights[name].get(key, torch.zeros_like(references[key])) for key in all_keys} + if name == "V0": + return ChangeTrackingDict(values) + return MathDict(values) variables = LazyVariableDict({name: value if value is not None else 0.0 for name, value in F.items()}) for alias, target in mapping.items(): diff --git a/tests/test_unified_visitor_regressions.py b/tests/test_unified_visitor_regressions.py new file mode 100644 index 0000000..dd9d2d1 --- /dev/null +++ b/tests/test_unified_visitor_regressions.py @@ -0,0 +1,90 @@ +import importlib +import os + +import pytest +import torch +from PIL import ImageFont + +from more_math.Parser.UnifiedMathVisitor import ChangeTrackingDict, MathDict, UnifiedMathVisitor +from more_math.helper_functions import parse_expr + + +def _visit(expression, variables): + return UnifiedMathVisitor(variables).visit(parse_expr(expression)) + + +def test_tracked_dict_arithmetic_preserves_original_baseline(): + values = { + "changed": torch.tensor([1.0, 2.0]), + "unchanged": torch.tensor([3.0, 4.0]), + } + result = _visit("V0 * 2", {"V0": ChangeTrackingDict(values)}) + + assert isinstance(result, ChangeTrackingDict) + assert result.get_changed_keys() == {"changed", "unchanged"} + assert torch.equal(result._initial_values["changed"], values["changed"]) + assert torch.equal(result["changed"], torch.tensor([2.0, 4.0])) + + +def test_tracked_dict_mapping_uses_v0_baseline_and_untracked_v1(): + v0 = ChangeTrackingDict({"weight": torch.tensor([1.0, 2.0])}) + v1 = MathDict({"weight": torch.tensor([3.0, 4.0])}) + + result = _visit("V0 + V1", {"V0": v0, "V1": v1}) + v1_result = _visit("V1", {"V0": v0, "V1": v1}) + + assert result.get_changed_keys() == {"weight"} + assert isinstance(v1_result, MathDict) + assert not isinstance(v1_result, ChangeTrackingDict) + + +@pytest.mark.parametrize( + "dtype_name", + ("float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz"), +) +def test_float8_operations_preserve_dtype(dtype_name): + dtype = getattr(torch, dtype_name, None) + if dtype is None: + pytest.skip(f"{dtype_name} is unavailable in this PyTorch build") + + values = torch.tensor([1.0, 2.0], dtype=dtype) + tensor_result = _visit("x * 2", {"x": values}) + unary_result = _visit("-x", {"x": values}) + dict_result = _visit("V0 * 2", {"V0": ChangeTrackingDict({"weight": values})}) + + assert tensor_result.dtype == dtype + assert unary_result.dtype == dtype + assert dict_result["weight"].dtype == dtype + assert torch.equal(tensor_result.float(), torch.tensor([2.0, 4.0])) + assert torch.equal(unary_result.float(), torch.tensor([-1.0, -2.0])) + assert torch.equal(dict_result["weight"].float(), torch.tensor([2.0, 4.0])) + + +@pytest.mark.parametrize("call", ("f()", "f(1, 2)")) +def test_lambda_calls_require_exact_argument_count(call): + with pytest.raises(ValueError, match="Lambda expects 1 arguments"): + _visit(f"f = (x) -> x; {call}", {}) + + +def test_text_image_size_expression_is_evaluated_once(monkeypatch): + visitor_module = importlib.import_module("more_math.Parser.UnifiedMathVisitor") + default_font = ImageFont.load_default() + original_exists = os.path.exists + monkeypatch.setattr( + visitor_module.ImageFont, + "truetype", + lambda *args, **kwargs: default_font, + ) + monkeypatch.setattr( + visitor_module.os.path, + "exists", + lambda path: path == "mock-font" or original_exists(path), + ) + + values = MathDict() + _visit( + 'text_image("x", "mock-font", count(add_key(d, count(d), 1)))', + {"d": values}, + ) + + assert len(values) == 1