From fc3f9351ca85511b93507ef30efb1e114baf5aa3 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Thu, 17 Sep 2026 17:29:53 +0200 Subject: [PATCH] AI: fix result not changing when overwriting input value in dictionary mode --- more_math/modelLikeCommon.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/more_math/modelLikeCommon.py b/more_math/modelLikeCommon.py index 468f728..57ed220 100644 --- a/more_math/modelLikeCommon.py +++ b/more_math/modelLikeCommon.py @@ -297,12 +297,20 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=[], us variables = variables | generate_dim_variables(ref_tensor) # Execute math + # Store original V0 before visiting, since visitor may modify it + original_v0 = variables.get("V0") + if original_v0 is not None: + # If it's a tensor, clone it; if it's a dict, copy its items + if isinstance(original_v0, torch.Tensor): + original_v0 = original_v0.clone() + elif isinstance(original_v0, dict): + original_v0 = MathDict({k: (v.clone() if isinstance(v, torch.Tensor) else v) for k, v in original_v0.items()}) visitor = UnifiedMathVisitor(variables, ref_tensor.shape,state_storage=stack) res = visitor.visit(tree) res = as_tensor(res, ref_tensor.shape) - original = variables.get("V0") + original = original_v0 if original is None: original = torch.zeros_like(res) @@ -436,6 +444,9 @@ def calculate_patches_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=[], u variables[f"{name}_d"] = make_delta_dict visitor_device = compute_device if compute_device is not None else next(iter(references.values())).device + # Store a copy of V0 BEFORE visiting, since the visitor may modify it + v0_dict = variables.get("V0", MathDict({key: torch.zeros_like(reference) for key, reference in references.items()})) + originals = MathDict({k: (v.clone() if isinstance(v, torch.Tensor) else v) for k, v in v0_dict.items()}) result = UnifiedMathVisitor(variables, device=visitor_device, state_storage=stack).visit(tree) if not isinstance(result, MathDict): if isinstance(result, (float, int)): @@ -445,8 +456,6 @@ def calculate_patches_dict_mode(Expr, V, F, pbar=None, mapping=None, stack=[], u if tuple(result) != tuple(all_keys): raise ValueError("Dictionary mode expressions must preserve all model keys") - - originals = variables.get("V0", MathDict({key: torch.zeros_like(reference) for key, reference in references.items()})) patches = {} for key, value in result.items(): if not torch.is_tensor(value):