chore: 🤖 fix pylint

yisongfu · yisongfu · commit aeb74ac1c8a5 · 2025-11-03T09:43:20.000+08:00
diff --git a/src/basicts/data/base_dataset.py b/src/basicts/data/base_dataset.py
@@ -1,8 +1,6 @@
 from typing import Union
 
 import numpy as np
-
-
 from basicts.utils.constants import BasicTSMode
 from torch.utils.data import Dataset
 
diff --git a/src/basicts/data/blast.py b/src/basicts/data/blast.py
@@ -4,6 +4,7 @@
 
 import numpy as np
 from basicts.utils.constants import BasicTSMode
+
 from .base_dataset import BasicTSDataset
 
 
diff --git a/src/basicts/data/tsf_dataset.py b/src/basicts/data/tsf_dataset.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 from basicts.utils.constants import BasicTSMode
+
 from .base_dataset import BasicTSDataset
 
 
diff --git a/src/basicts/data/tsi_dataset.py b/src/basicts/data/tsi_dataset.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 from basicts.utils.constants import BasicTSMode
+
 from .base_dataset import BasicTSDataset
 
 
diff --git a/src/basicts/data/uea_dataset.py b/src/basicts/data/uea_dataset.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 from basicts.utils import BasicTSMode
+
 from .base_dataset import BasicTSDataset
 
 
@@ -93,5 +94,6 @@ def __len__(self) -> int:
         """
         return self.inputs.shape[0]
 
+    @property
     def data(self) -> np.ndarray:
         return self.inputs
diff --git a/src/basicts/models/Autoformer/arch/autoformer_arch.py b/src/basicts/models/Autoformer/arch/autoformer_arch.py
@@ -7,7 +7,9 @@
 from torch import nn
 
 from ..config.autoformer_config import AutoformerConfig
-from .layers import AutoformerDecoder, AutoformerDecoderLayer, AutoformerEncoderLayer
+from .layers import (AutoformerDecoder, AutoformerDecoderLayer,
+                     AutoformerEncoderLayer)
+
 
 class Autoformer(nn.Module):
     """
diff --git a/src/basicts/models/Autoformer/arch/layers.py b/src/basicts/models/Autoformer/arch/layers.py
@@ -1,9 +1,8 @@
 from typing import Any, Callable, Optional, Tuple
 
 import torch
-from torch import nn
-
 from basicts.modules.transformer.utils import build_layer
+from torch import nn
 
 from ..config.autoformer_config import AutoformerConfig
 
diff --git a/src/basicts/models/DUET/arch/duet_arch.py b/src/basicts/models/DUET/arch/duet_arch.py
@@ -1,12 +1,12 @@
-import torch
-from typing import Dict, Any, List
-
-from torch import nn
-from einops import rearrange
+from typing import Any, Dict, List
 
-from basicts.modules.transformer import Encoder, EncoderLayer, MultiHeadAttention
+import torch
 from basicts.modules.mlps import MLPLayer
+from basicts.modules.transformer import (Encoder, EncoderLayer,
+                                         MultiHeadAttention)
 from basicts.runners.callback import AddAuxiliaryLoss
+from einops import rearrange
+from torch import nn
 
 from ..config.duet_config import DUETConfig
 from .linear_extractor_cluster import LinearExtractorCluster
diff --git a/src/basicts/models/DUET/arch/linear_extractor_cluster.py b/src/basicts/models/DUET/arch/linear_extractor_cluster.py
@@ -1,11 +1,10 @@
 import torch
 import torch.nn as nn
-from torch.distributions.normal import Normal
-from einops import rearrange
-
-from basicts.modules.norm import RevIN
-from basicts.modules.mlps import MLPLayer
 from basicts.models.DLinear import DLinear, DLinearConfig
+from basicts.modules.mlps import MLPLayer
+from basicts.modules.norm import RevIN
+from einops import rearrange
+from torch.distributions.normal import Normal
 
 from ..config.duet_config import DUETConfig
 
diff --git a/src/basicts/models/DUET/arch/mahalanobis_mask.py b/src/basicts/models/DUET/arch/mahalanobis_mask.py
@@ -1,7 +1,7 @@
 import torch
 import torch.nn.functional as F
-from torch import nn
 from einops import rearrange
+from torch import nn
 
 
 class MahalanobisMask(nn.Module):
diff --git a/src/basicts/models/LightTS/arch/lightts_arch.py b/src/basicts/models/LightTS/arch/lightts_arch.py
@@ -1,7 +1,6 @@
 import torch
-from torch import nn
-
 from basicts.modules import MLPLayer
+from torch import nn
 
 from ..config.lightts_config import LightTSConfig
 
diff --git a/src/basicts/models/MTSMixer/arch/mtsmixer_arch.py b/src/basicts/models/MTSMixer/arch/mtsmixer_arch.py
@@ -1,11 +1,9 @@
 import torch
 import torch.nn as nn
-
 from basicts.modules.norm import RevIN
 
 from ..config.mtsmixer_config import MTSMixerConfig
-
-from .mtsmixer_layers import MixerLayer, ChannelProjection
+from .mtsmixer_layers import ChannelProjection, MixerLayer
 
 
 class MTSMixer(nn.Module):
diff --git a/src/basicts/models/MTSMixer/arch/mtsmixer_layers.py b/src/basicts/models/MTSMixer/arch/mtsmixer_layers.py
@@ -1,6 +1,7 @@
 import torch
-from torch import nn
 from basicts.modules import MLPLayer
+from torch import nn
+
 from ..config.mtsmixer_config import MTSMixerConfig
 
 
diff --git a/src/basicts/models/TimeKAN/arch/timekan_arch.py b/src/basicts/models/TimeKAN/arch/timekan_arch.py
@@ -1,8 +1,7 @@
 import torch
-from torch import nn
-
 from basicts.modules.embed import FeatureEmbedding
 from basicts.modules.norm import RevIN
+from torch import nn
 
 from ..config.timekan_config import TimeKANConfig
 from .timekan_layers import FrequencyDecompLayer, FrequencyMixingLayer
diff --git a/src/basicts/models/TimeXer/arch/layers.py b/src/basicts/models/TimeXer/arch/layers.py
@@ -1,9 +1,8 @@
 from typing import Callable, Tuple
 
 import torch
-from torch import nn
-
 from basicts.modules.transformer import Seq2SeqDecoderLayer
+from torch import nn
 
 
 class FlattenHead(nn.Module):
diff --git a/src/basicts/models/TimeXer/arch/timexer_arch.py b/src/basicts/models/TimeXer/arch/timexer_arch.py
@@ -1,10 +1,9 @@
 import torch
-from torch import nn
-
 from basicts.modules.embed import PatchEmbedding, SequenceEmbedding
 from basicts.modules.mlps import MLPLayer
 from basicts.modules.norm import RevIN
 from basicts.modules.transformer import MultiHeadAttention, Seq2SeqDecoder
+from torch import nn
 
 from ..config.timexer_config import TimeXerConfig
 from .layers import FlattenHead, TimeXerEncoderLayer
diff --git a/src/basicts/modules/__init__.py b/src/basicts/modules/__init__.py
@@ -1,8 +1,8 @@
 from basicts.modules.activations import ACT2FN
 from basicts.modules.decomposition import (DFTDecomposition, MovingAverage,
-                                           MovingAverageDecomposition)
+                                           MovingAverageDecomposition,
+                                           MultiMovingAverageDecomposition)
 from basicts.modules.mlps import MLPLayer, ResMLPLayer
-from basicts.modules.decomposition import DFTDecomposition, MovingAverage, MovingAverageDecomposition, MultiMovingAverageDecomposition
 
 __ALL__ = [
     "ACT2FN",
diff --git a/src/basicts/modules/decomposition.py b/src/basicts/modules/decomposition.py
@@ -1,4 +1,4 @@
-from typing import Tuple, Sequence
+from typing import Sequence, Tuple
 
 import torch
 from torch import nn
diff --git a/src/basicts/modules/transformer/decoder.py b/src/basicts/modules/transformer/decoder.py
@@ -1,4 +1,4 @@
-from typing import Callable, Literal, Optional, Tuple
+from typing import Callable, Literal, Optional, Tuple, Union
 
 import torch
 from torch import nn
@@ -17,7 +17,7 @@ def __init__(
         self,
         self_attn: nn.Module,
         ffn_layer: nn.Module,
-        layer_norm: Callable | Tuple[Callable, dict],
+        layer_norm: Union[Callable, Tuple[Callable, dict]],
         norm_position: Literal["pre", "post", "sandwich"] = "pre",
     ):
         super().__init__()
@@ -163,7 +163,7 @@ def __init__(
         self_attn: nn.Module,
         cross_attn: nn.Module,
         ffn_layer: nn.Module,
-        layer_norm: Callable | Tuple[Callable, dict],
+        layer_norm: Union[Callable, Tuple[Callable, dict]],
         norm_position: Literal["pre", "post", "sandwich"] = "pre",
     ):
         super().__init__()
diff --git a/src/basicts/modules/transformer/encoder.py b/src/basicts/modules/transformer/encoder.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, List, Literal, Optional, Tuple
+from typing import Any, Callable, List, Literal, Optional, Tuple, Union
 
 import torch
 from torch import nn
@@ -16,7 +16,7 @@ def __init__(
         self,
         self_attn: nn.Module,
         ffn_layer: nn.Module,
-        layer_norm: Callable | Tuple[Callable, Any],
+        layer_norm: Union[Callable, Tuple[Callable, Any]],
         norm_position: Literal["pre", "post", "sandwich"] = "pre"
     ):
         super().__init__()
diff --git a/src/basicts/modules/transformer/utils.py b/src/basicts/modules/transformer/utils.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, Optional, Tuple, Union
 
 import torch
 from torch import nn
@@ -53,7 +53,7 @@ def prepare_causal_attention_mask(
     return combined_mask
 
 def build_layer(
-    layer: Callable | Tuple[Callable, Any],
+    layer: Union[Callable, Tuple[Callable, Any]],
 ) -> nn.Module:
     """Build layer.
 
diff --git a/src/basicts/runners/basicts_runner.py b/src/basicts/runners/basicts_runner.py
@@ -131,7 +131,7 @@ def __init__(self, cfg: "BasicTSConfig") -> None:
         self.lr_scheduler = None
 
         # automatic mixed precision (amp)
-        model_dtype: str | torch.dtype = cfg.get("model_dtype", "float32")
+        model_dtype: Union[str, torch.dtype] = cfg.get("model_dtype", "float32")
         if isinstance(model_dtype, str):
             self.ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[model_dtype]
         else:
diff --git a/src/basicts/runners/builder.py b/src/basicts/runners/builder.py
@@ -98,9 +98,9 @@ def _build_data_loader(cfg: "BasicTSConfig", mode: BasicTSMode, logger: Logger):
         ) if torch.distributed.is_initialized() and mode == BasicTSMode.TRAIN else None
 
         shuffle = False if torch.distributed.is_initialized() and mode == BasicTSMode.TRAIN\
-              else cfg.get(f'{mode}_data_shuffle', False)
+              else cfg.get(f"{mode}_data_shuffle", False)
 
-        return (DataLoaderX if cfg.get(f"{mode}_data_prefetch") else DataLoader)(
+        return (DataLoaderX if cfg.get(f"{mode}_data_prefetch", False) else DataLoader)(
             dataset,
             collate_fn=cfg.get(f"{mode}_data_collate_fn", None),
             batch_size=cfg.get(f"{mode}_batch_size", 1),
diff --git a/src/basicts/runners/callback/__init__.py b/src/basicts/runners/callback/__init__.py
@@ -7,7 +7,6 @@
 from .no_bp import NoBP
 from .selective_learning import SelectiveLearning
 
-
 __ALL__ = [
     'AddAuxiliaryLoss',
     'BasicTSCallback',
diff --git a/src/basicts/runners/callback/grad_accumulation.py b/src/basicts/runners/callback/grad_accumulation.py
@@ -23,8 +23,8 @@ def on_train_start(self, runner: "BasicTSRunner"):
         world_size = torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1
         batch_size = runner.cfg.train_batch_size
         effective_batch_size = batch_size * world_size * self.accumulation_steps
-        runner.logger.info(f'Use gradient accumulation with step {self.accumulation_steps}.')
-        runner.logger.info(f'Effective batch size: {effective_batch_size}.')
+        runner.logger.info(f"Use gradient accumulation with step {self.accumulation_steps}.")
+        runner.logger.info(f"Effective batch size: {effective_batch_size}.")
 
     def on_backward(self, runner: "BasicTSRunner", loss: torch.Tensor):
         runner.should_optimizer_step = (self.current_steps + 1) % self.accumulation_steps == 0
