Use both torch.fx.wrap and autowrap fns list to avoid some issues

Author: rwightman

timm/layers/_fx.py

@@ -46,17 +46,10 @@ def get_notrace_modules():
 # Functions we want to autowrap (treat them as leaves)
 _autowrap_functions = set()
 
-try:
-    # pass through to torch.fx.wrap when possible, works in some cases our old mechanism doesn't
-    register_notrace_function = torch.fx.wrap  # exists in modern PyTorch
-except AttributeError:
-    # old Torch
-    def register_notrace_function(name_or_fn):
-        if callable(name_or_fn):
-            _autowrap_functions.add(name_or_fn)
-            return name_or_fn
-        _autowrap_functions.add(name_or_fn)
-        return name_or_fn
+
+def register_notrace_function(name_or_fn):
+    _autowrap_functions.add(name_or_fn)
+    return name_or_fn
 
 
 def is_notrace_function(func: Callable):

timm/layers/attention.py

@@ -9,6 +9,7 @@
 from .pos_embed_sincos import apply_rot_embed_cat
 
 
+@torch.fx.wrap
 @register_notrace_function
 def maybe_add_mask(scores: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
     return scores if attn_mask is None else scores + attn_mask

timm/layers/attention_pool2d.py

@@ -12,16 +12,12 @@
 import torch
 import torch.nn as nn
 
-from ._fx import register_notrace_function
 from .config import use_fused_attn
 from .helpers import to_2tuple
 from .pos_embed import resample_abs_pos_embed
 from .pos_embed_sincos import apply_rot_embed, RotaryEmbedding
 from .weight_init import trunc_normal_
 
-# have to register again for some reason
-register_notrace_function(resample_abs_pos_embed)
-
 
 class RotAttentionPool2d(nn.Module):
     """ Attention based 2D feature pooling w/ rotary (relative) pos embedding.

timm/layers/pos_embed.py

@@ -14,6 +14,7 @@
 _logger = logging.getLogger(__name__)
 
 
+@torch.fx.wrap
 @register_notrace_function
 def resample_abs_pos_embed(
         posemb: torch.Tensor,
@@ -58,6 +59,7 @@ def resample_abs_pos_embed(
     return posemb
 
 
+@torch.fx.wrap
 @register_notrace_function
 def resample_abs_pos_embed_nhwc(
         posemb: torch.Tensor,

timm/layers/pos_embed_sincos.py

@@ -512,6 +512,7 @@ def init_random_2d_freqs(
     return torch.stack([fx, fy], dim=0)
 
 
+@torch.fx.wrap
 @register_notrace_function
 def get_mixed_freqs(
         freqs: torch.Tensor,
@@ -584,8 +585,6 @@ def __init__(
         )  # (2, depth, num_heads, head_dim//2)
         self.freqs = nn.Parameter(freqs)
 
-
-
     def get_embed(self, shape: Optional[List[int]] = None) -> torch.Tensor:
         """Generate rotary embeddings for the given spatial shape.