Add Torch flex attention endpoint compatibility shim

lagunavision/backbones/hf_causal.py

@@ -83,6 +83,17 @@ def _install_torch_compat_shims() -> None:
     except Exception:
         pass
 
+    try:
+        import torch.nn.attention.flex_attention as flex_attention
+
+        if not hasattr(flex_attention, "AuxRequest"):
+            class AuxRequest:  # noqa: N801 - mirrors PyTorch internal name
+                pass
+
+            flex_attention.AuxRequest = AuxRequest
+    except Exception:
+        pass
+
 
 def _install_huggingface_hub_dataclasses_shim() -> None:
     """Provide compatibility for Laguna remote code on HF Endpoint runtimes."""