Update

[ghstack-poisoned]
pytorch · bdhirsh · Apr 18, 2025 · Apr 25, 2025 · Apr 25, 2025 · Apr 28, 2025
commit 398bc2ec7dd9750cab6b809532fec5a4f3ea084d
diff --git a/torch/_higher_order_ops/flex_attention.py b/torch/_higher_order_ops/flex_attention.py
@@ -401,7 +401,6 @@ def flex_attention_functionalize(
     """
     from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
 
-
     if has_user_subclass(
         (
             query,
@@ -479,7 +478,6 @@ def flex_attention_fake_impl(
     score_mod_other_buffers: tuple = (),
     mask_mod_other_buffers: tuple = (),
 ) -> tuple[torch.Tensor, torch.Tensor]:
-
     if has_user_subclass(
         (
             query,
@@ -1215,7 +1213,6 @@ def flex_attention_backward_fake_tensor_mode(
 ) -> tuple[
     torch.Tensor, torch.Tensor, torch.Tensor, tuple[Optional[torch.Tensor], ...]
 ]:
-
     if has_user_subclass(
         (
             query,

diff --git a/torch/_higher_order_ops/utils.py b/torch/_higher_order_ops/utils.py
@@ -1102,14 +1102,14 @@ def wrapped_fn(*flat_args):
 
 def has_user_subclass(args, allowed_subclasses):
     """Check if any tensor arguments are user subclasses.
-    
-    This is used to determine if tensor subclasses should get a chance to run 
+
+    This is used to determine if tensor subclasses should get a chance to run
     their own implementation first before falling back to the default implementation.
-    
+
     Args:
         args: Arguments to check (will be flattened with pytree)
         allowed_subclasses: Tuple of allowed subclass types
-        
+
     Returns:
         True if user tensor subclasses are found, False otherwise
     """