fix: resolve NaN/Inf issues in 17 illegal Torch samples

- Add -inf to -1e6 replacement logic in apply_templates for newly generated models
- Add runtime replacement logic in load_class_from_file for existing models
- Fix NaN issues in masked_fill and torch.full calls that use -inf
- Ensure all 17 samples pass test_compiler with both nope and inductor backends
- No manual modification of auto-generated model.py files required

Author: Dayuxiaoshui

graph_net/torch/test_compiler.py

@@ -61,6 +61,47 @@ def load_class_from_file(
 
     with open(file_path, "r", encoding="utf-8") as f:
         model_code = f.read()
+    # Replace -inf with -1e6 in masked_fill and torch.full to prevent NaN issues
+    # This applies the same fix as apply_templates for already-generated model.py files
+    import re
+    # Pattern for masked_fill(..., -inf)
+    model_code = re.sub(
+        r"(masked_fill\([^,)]+,\s*)-inf(\s*\))", r"\1-1e6\2", model_code
+    )
+    # For torch.full, use a context-aware replacement
+    # Find torch.full(...) blocks and replace -inf within them
+    # Use a balanced bracket matcher approach
+    parts = []
+    i = 0
+    while i < len(model_code):
+        if model_code[i:].startswith('torch.full('):
+            # Find the matching closing parenthesis
+            depth = 0
+            start = i
+            j = i + len('torch.full(')
+            while j < len(model_code):
+                if model_code[j] == '(':
+                    depth += 1
+                elif model_code[j] == ')':
+                    if depth == 0:
+                        # Found the matching closing paren
+                        full_block = model_code[start:j+1]
+                        # Replace -inf with -1e6 in this block
+                        full_block = full_block.replace('-inf', '-1e6')
+                        parts.append(full_block)
+                        i = j + 1
+                        break
+                    depth -= 1
+                j += 1
+            else:
+                # Didn't find closing paren, just append rest
+                parts.append(model_code[i:])
+                break
+        else:
+            parts.append(model_code[i])
+            i += 1
+    if parts:
+        model_code = ''.join(parts)
     model_code = utils.modify_code_by_device(model_code, device)
     spec = importlib.util.spec_from_loader(module_name, loader=None)
     module = importlib.util.module_from_spec(spec)

graph_net/torch/utils.py

@@ -17,10 +17,20 @@
 def apply_templates(forward_code: str) -> str:
     tab = "    "
     forward_code = f"\n{tab}".join(forward_code.split("\n"))
+    # Replace -inf with -1e6 in masked_fill and torch.full to prevent NaN issues
+    # This is specifically for cases where -inf can cause problems (e.g., before sigmoid)
+    # Pattern for masked_fill(..., -inf)
+    forward_code = re.sub(
+        r"(masked_fill\([^,)]+,\s*)-inf(\s*\))", r"\1-1e6\2", forward_code
+    )
+    # Pattern for torch.full(..., -inf, ...)
+    forward_code = re.sub(
+        r"(torch\.full\([^,)]+,\s*)-inf(\s*[,)])", r"\1-1e6\2", forward_code
+    )
     imports = "import torch"
     if "device" in forward_code:
         imports += "\n\nfrom torch import device"
-    if "inf" in forward_code:
+    if "inf" in forward_code or "-1e6" in forward_code:
         imports += "\n\nfrom torch import inf"
     return f"{imports}\n\nclass GraphModule(torch.nn.Module):\n{tab}{forward_code}"