[flex attention] change "==" to "is" in inspect parameter comparison (pytorch#165003)

williamwen42 · Chao1Han · commit 476910f91301 · 2025-10-21T15:38:02.000+08:00
Patch for pytorch#164760. This doesn't actually fix the underlying torch function issue though. Explanation: `is` is traced differently compared to `__eq__`, so we end up avoiding the issue where we attempt to evaluate `torch.eq(tensor, inspect._empty)` in the first place. Pull Request resolved: pytorch#165003 Approved by: https://github.com/mlazos
diff --git a/torch/nn/attention/flex_attention.py b/torch/nn/attention/flex_attention.py
@@ -261,7 +261,7 @@ def _get_mod_type(fn: Callable) -> _ModificationType:
     num_positional_args = sum(
         1
         for param in inspect.signature(fn).parameters.values()
-        if param.default == inspect.Parameter.empty
+        if param.default is inspect.Parameter.empty
     )
     assert num_positional_args == 5 or num_positional_args == 4
     if num_positional_args == 5:

Original file line number	Diff line number	Diff line change
`@@ -261,7 +261,7 @@ def _get_mod_type(fn: Callable) -> _ModificationType:`
`261`	`261`	`num_positional_args = sum(`
`262`	`262`	`1`
`263`	`263`	`for param in inspect.signature(fn).parameters.values()`
`264`		`- if param.default == inspect.Parameter.empty`
	`264`	`+ if param.default is inspect.Parameter.empty`
`265`	`265`	`)`
`266`	`266`	`assert num_positional_args == 5 or num_positional_args == 4`
`267`	`267`	`if num_positional_args == 5:`