[Fix] Stringify torch.device annotations for Paddle

fla/layers/mamba3.py

@@ -81,7 +81,7 @@ def __init__(
         use_bias: bool = False,
         norm_eps: float = 1e-5,
         layer_idx: int | None = None,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> Mamba3:
         super().__init__()
@@ -449,7 +449,7 @@ def allocate_inference_cache(
         self,
         batch_size: int,
         max_seqlen: int,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ):
         device = device or self.in_proj.weight.device

fla/layers/yoco.py

@@ -38,7 +38,7 @@ def __init__(
         scale_base: float | None = None,
         interleaved: bool = False,
         pos_idx_in_fp32: bool = True,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         rope_inv_freq: str = 'fla',
     ):
         self.rope_inv_freq = rope_inv_freq

fla/modules/conv/short_conv.py

@@ -62,7 +62,7 @@ def __init__(
         bias: bool = False,
         activation: str | None = 'silu',
         backend: str | None = 'triton',
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
         **kwargs,
     ):

fla/modules/feature_map.py

@@ -20,7 +20,7 @@
 
 
 @functools.cache
-def _triu_indices(n: int, offset: int, device: torch.device) -> torch.Tensor:
+def _triu_indices(n: int, offset: int, device: "torch.device") -> torch.Tensor:
     # cache the upper-triangular gather indices per (size, offset, device) to avoid rebuilding
     # them and copying host -> device on every forward
     return torch.triu_indices(n, n, offset, device=device)

fla/modules/fused_norm_gate.py

@@ -933,7 +933,7 @@ def __init__(
         bias: bool = False,
         activation: str = "swish",
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> FusedLayerNormGated:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -999,7 +999,7 @@ def __init__(
         elementwise_affine: bool = True,
         eps: float = 1e-5,
         activation: str = "swish",
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> FusedRMSNormGated:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -1062,7 +1062,7 @@ def __init__(
         elementwise_affine: bool = True,
         bias: bool = False,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> FusedLayerNormSwishGate:
         super().__init__(
@@ -1081,7 +1081,7 @@ def __init__(
         hidden_size: int,
         elementwise_affine: bool = True,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> FusedRMSNormSwishGate:
         super().__init__(
@@ -1099,7 +1099,7 @@ def __init__(
         hidden_size: int,
         elementwise_affine: bool = True,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> FusedLayerNormGatedLinear:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -1159,7 +1159,7 @@ def __init__(
         hidden_size: int,
         elementwise_affine: bool = True,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> FusedLayerNormSwishGateLinear:
         super().__init__(
@@ -1177,7 +1177,7 @@ def __init__(
         hidden_size,
         elementwise_affine: bool = True,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> FusedRMSNormGatedLinear:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -1236,7 +1236,7 @@ def __init__(
         hidden_size: int,
         elementwise_affine: bool = True,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> FusedRMSNormSwishGateLinear:
         super().__init__(

fla/modules/layernorm.py

@@ -980,7 +980,7 @@ def __init__(
         elementwise_affine: bool = True,
         bias: bool = False,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> LayerNorm:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -1035,7 +1035,7 @@ def __init__(
         bias: bool = False,
         eps: float = 1e-5,
         is_rms_norm: bool = False,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> GroupNorm:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -1097,7 +1097,7 @@ def __init__(
         elementwise_affine: bool = True,
         bias: bool = False,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> RMSNorm:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -1255,7 +1255,7 @@ def __init__(
         elementwise_affine: bool = True,
         bias: bool = False,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> LayerNormLinear:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -1313,7 +1313,7 @@ def __init__(
         bias: bool = False,
         eps: float = 1e-5,
         is_rms_norm: bool = False,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> GroupNormLinear:
         factory_kwargs = {"device": device, "dtype": dtype}
@@ -1377,7 +1377,7 @@ def __init__(
         elementwise_affine: bool = True,
         bias: bool = False,
         eps: float = 1e-5,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ) -> RMSNormLinear:
         factory_kwargs = {"device": device, "dtype": dtype}

fla/modules/layernorm_gated.py

@@ -475,7 +475,7 @@ def __init__(
         eps: float = 1e-5,
         group_size: int | None = None,
         norm_before_gate: bool = True,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ):
         """If group_size is not None, we do GroupNorm with each group having group_size elements.
@@ -510,7 +510,7 @@ def __init__(
         eps: float = 1e-5,
         group_size: int | None = None,
         norm_before_gate: bool = False,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
         dtype: torch.dtype | None = None,
     ):
         """If group_size is not None, we do GroupNorm with each group having group_size elements.

fla/modules/rotary.py

@@ -342,7 +342,7 @@ def __init__(
         scale_base: float | None = None,
         interleaved: bool = False,
         pos_idx_in_fp32: bool = True,
-        device: torch.device | None = None,
+        device: "torch.device | None" = None,
     ):
         """
         interleaved:

fla/ops/attnres/backends/triton_ascend/fused.py

@@ -388,7 +388,7 @@ def _get_o_mix(
     logit: torch.Tensor,
     lse: torch.Tensor,
     scale: float,
-    device: torch.device,
+    device: "torch.device",
     o_pre: torch.Tensor | None = None,
     o_mix: torch.Tensor | None = None,
 ) -> torch.Tensor:

fla/ops/common/intracard_cp.py

@@ -297,7 +297,7 @@ def intracard_merge(
     num_non_first: int,
     merge_seq_offsets: list[int],
     merge_init_offsets: list[int],
-    device: torch.device,
+    device: "torch.device",
     initial_state: torch.Tensor | None = None,
     state_v_first: bool = False,
 ) -> tuple[torch.Tensor | None, int]:

fla/ops/gated_delta_rule/backends/triton_ascend/wy_fast.py

@@ -50,7 +50,7 @@ def _beta_npu_arg(beta: torch.Tensor, HV: int) -> tuple[torch.Tensor, bool]:
 
 
 def _t_npu_buf(
-    B: int, T: int, HV: int, *, dtype: torch.dtype, device: torch.device,
+    B: int, T: int, HV: int, *, dtype: torch.dtype, device: "torch.device",
 ) -> tuple[torch.Tensor, bool]:
     if HV == 1:
         return torch.empty(B, T, HV, dtype=dtype, device=device), False

fla/ops/utils/index.py

@@ -72,7 +72,7 @@ def prepare_split_cu_seqlens(
     split_size: int | None = None,
     cu_seqlens: torch.LongTensor | None = None,
     dtype: torch.dtype | None = torch.int32,
-    device: torch.device | None = torch.device('cpu'),
+    device: "torch.device | None" = torch.device('cpu'),
 ) -> torch.LongTensor:
     """Sub-split a (optionally packed) batch along the token axis.