nodes/mcDandy-more_math
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

AI: sstylistic and fixes of visitor

Browse Source
This commit is contained in:
mcDandy
2026-10-13 15:58:07 +02:00
parent 87f61da697
commit a925909e10
3 changed files with 187 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
more_math/Parser/UnifiedMathVisitor.py
+93 -63
View File
@@ -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
more_math/modelLikeCommon.py
+4 -3
View File
@@ -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():
tests/test_unified_visitor_regressions.py
+90
View File
@@ -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