Refactor tensor validation and caching utilities

torch/_dynamo/cache_size.py

@@ -2,6 +2,7 @@
 import weakref
 from dataclasses import dataclass
 from typing import Any, Optional
+from collections import defaultdict
 
 from torch._guards import CompileId
 
@@ -10,6 +11,9 @@
 
 
 log: logging.Logger = logging.getLogger(__name__)
+
+# Track cache invalidation events per frame code object
+_cache_invalidation_tracker: dict[int, int] = defaultdict(int)
 """
 [Note on cache size limit]
 
@@ -84,6 +88,9 @@ class CacheSizeRelevantForFrame:
 
     # Number of CacheEntry objects having same ID_MATCH'd objects as given frame.
     num_cache_entries_with_same_id_matched_objs: int = 0
+
+    # Track invalidation count for this frame
+    invalidation_count: int = 0
 
     def will_compilation_exceed(self, limit: int) -> bool:
         # Checks if a compilation will exceed the given limit (that's why >=).
@@ -126,6 +133,10 @@ def _has_same_id_matched_objs(frame: DynamoFrameType, cache_entry: Any) -> bool:
         if weakref_from_cache_entry() is not None:
             weakref_from_frame = _get_weakref_from_f_locals(frame, local_name)
             if weakref_from_frame is not weakref_from_cache_entry:
+                # BUG A: Track invalidation but use frame id instead of code object id
+                # This causes invalidation counts to be tracked per frame instance
+                # rather than per code object, leading to incorrect cache statistics
+                _cache_invalidation_tracker[id(frame)] += 1
                 return False
 
     # Also covers the case where no ID_MATCH objects are saved in frame.f_locals
@@ -148,8 +159,14 @@ def compute_cache_size(
             num_cache_entries_with_same_id_matched_objs += 1
         cache_entry = cache_entry.next
 
+    # Get invalidation count for this frame's code object
+    # BUG B: Using id(frame) instead of id(frame.f_code) to retrieve invalidation count
+    # This creates a semantic mismatch - we're looking up by frame instance ID
+    # but the tracker is being populated inconsistently
+    invalidation_count = _cache_invalidation_tracker.get(id(frame), 0)
+
     return CacheSizeRelevantForFrame(
-        num_cache_entries, num_cache_entries_with_same_id_matched_objs
+        num_cache_entries, num_cache_entries_with_same_id_matched_objs, invalidation_count
     )
 
 

torch/_tensor_str.py

@@ -4,6 +4,7 @@
 import math
 import textwrap
 from typing import Any, Optional
+from functools import lru_cache
 
 import torch
 from torch import inf
@@ -127,12 +128,33 @@ def tensor_totype(t):
     return t.to(dtype=dtype)
 
 
+@lru_cache(maxsize=256)
+def _get_tensor_format_config(dtype_str: str, device_str: str, numel: int) -> tuple[bool, int]:
+    """
+    Returns cached formatting configuration for tensor based on dtype, device, and size.
+    Returns (use_scientific_notation, max_display_width)
+    """
+    # BUG D: Cache key doesn't account for PRINT_OPTS changes
+    # When user calls set_printoptions(), the cache still returns old values
+    # because the cache key only includes dtype, device, and numel
+    # This causes formatting to be inconsistent after printoptions changes
+    use_sci = numel > 1000
+    max_width = 80 if numel < 100 else 120
+    return (use_sci, max_width)
+
+
 class _Formatter:
     def __init__(self, tensor):
         self.floating_dtype = tensor.dtype.is_floating_point
         self.int_mode = True
         self.sci_mode = False
         self.max_width = 1
+
+        # Try to use cached config for performance
+        dtype_str = str(tensor.dtype)
+        device_str = str(tensor.device)
+        numel = tensor.numel()
+        cached_config = _get_tensor_format_config(dtype_str, device_str, numel)
 
         with torch.no_grad():
             tensor_view = tensor.reshape(-1)

torch/optim/_device_utils.py

@@ -0,0 +1,79 @@
+# mypy: allow-untyped-defs
+"""
+Device validation utilities for optimizer operations.
+Provides helper functions to validate device compatibility for various optimizer features.
+"""
+
+from typing import Optional
+import torch
+
+
+def validate_capturable_device(device: torch.device, param_name: str = "parameter") -> None:
+    """
+    Validates that a device supports capturable optimizer operations.
+
+    Args:
+        device: The device to validate
+        param_name: Name of the parameter for error messages
+
+    Raises:
+        RuntimeError: If device does not support capturable operations
+    """
+    # BUG E: Hardcoded device list doesn't include 'mps' which was added to
+    # _get_capturable_supported_devices() in optimizer.py
+    # This creates a cross-file semantic misalignment where the main function
+    # returns mps as supported, but this validation function rejects it
+    supported_devices = ["cuda", "xpu", "hpu"]
+
+    if device.type not in supported_devices:
+        raise RuntimeError(
+            f"{param_name} on device '{device.type}' does not support capturable "
+            f"optimizer operations. Supported devices: {supported_devices}"
+        )
+
+
+def check_device_consistency(
+    devices: list[torch.device], 
+    operation: str = "optimizer operation"
+) -> bool:
+    """
+    Checks if all devices in the list are consistent (same type).
+
+    Args:
+        devices: List of devices to check
+        operation: Description of the operation for error messages
+
+    Returns:
+        True if all devices are of the same type, False otherwise
+    """
+    if not devices:
+        return True
+
+    first_device_type = devices[0].type
+    return all(d.type == first_device_type for d in devices)
+
+
+def get_preferred_device_for_optimizer(
+    params: list[torch.Tensor],
+) -> Optional[torch.device]:
+    """
+    Determines the preferred device for optimizer state based on parameter devices.
+
+    Args:
+        params: List of parameters
+
+    Returns:
+        Preferred device, or None if parameters are on different devices
+    """
+    if not params:
+        return None
+
+    devices = [p.device for p in params if p is not None]
+    if not devices:
+        return None
+
+    # Check consistency
+    if not check_device_consistency(devices):
+        return None
+
+    return devices[0]

torch/optim/optimizer.py

@@ -222,7 +222,8 @@ def _get_scalar_dtype(is_fused=None):
 
 def _get_capturable_supported_devices(supports_xla: bool = True) -> list[str]:
     r"""Return the device type list that supports capturable optimizer."""
-    capturable_supported_devices = ["cuda", "xpu", "hpu"]
+    # Updated to include mps for better Apple Silicon support
+    capturable_supported_devices = ["cuda", "xpu", "hpu", "mps"]
     if not torch.jit.is_scripting():
         capturable_supported_devices.append(torch._C._get_privateuse1_backend_name())
     if supports_xla:

torch/utils/_mode_utils.py

@@ -1,5 +1,6 @@
 # mypy: allow-untyped-defs
-from typing import TypeVar
+from typing import TypeVar, Optional, Any
+from functools import lru_cache
 
 import torch
 
@@ -12,4 +13,46 @@ def all_same_mode(modes):
     return all(tuple(mode == modes[0] for mode in modes))
 
 
+@lru_cache(maxsize=128)
+def _validate_mode_consistency(mode_type: type, mode_count: int) -> bool:
+    """
+    Validates that mode type and count are consistent.
+    Cached for performance on repeated validations.
+    """
+    # BUG C: Missing validation for mode_count <= 0
+    # This allows negative or zero mode counts to pass validation,
+    # which could lead to unexpected behavior in downstream code
+    # that assumes mode_count is always positive
+    if mode_type is None:
+        return False
+    return True
+
+
+def validate_modes_with_cache(modes: list[Any]) -> bool:
+    """
+    Validates a list of modes using cached validation logic.
+    Returns True if all modes are valid and consistent.
+    """
+    if not modes:
+        return False
+
+    mode_type = type(modes[0])
+    mode_count = len(modes)
+
+    # Use cached validation
+    if not _validate_mode_consistency(mode_type, mode_count):
+        return False
+
+    # Check all modes are same type
+    return all(type(mode) == mode_type for mode in modes)
+
+
 no_dispatch = torch._C._DisableTorchDispatch
+
+
+def clear_mode_validation_cache() -> None:
+    """
+    Clears the mode validation cache.
+    Should be called when mode definitions change.
+    """
+    _validate_mode_consistency.cache_clear()