format and lint

Signed-off-by: peter szemraj <peterszemraj@gmail.com>

Author: pszemraj

samba_pytorch/config.py

@@ -7,11 +7,10 @@
 from dataclasses import dataclass
 from typing import Any, Literal, Optional, Type
 
-import torch
-from typing_extensions import Self
-
 import lit_gpt.model
+import torch
 from lit_gpt.utils import find_multiple
+from typing_extensions import Self
 
 
 @dataclass

samba_pytorch/modules/fused_rotary_embedding.py

@@ -3,12 +3,10 @@
 
 # Copyright (c) 2023, Tri Dao.
 
-import math
-from typing import Optional, Tuple
 
 import rotary_emb
 import torch
-from einops import rearrange, repeat
+from einops import rearrange
 
 
 class ApplyRotaryEmb(torch.autograd.Function):

samba_pytorch/modules/gla.py

@@ -13,15 +13,13 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache
-
 from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
 from fla.ops.gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
 
 
 class GatedLinearAttention(nn.Module):
-
     def __init__(
         self,
         mode: str = "fused_chunk",

samba_pytorch/modules/mamba_simple.py

@@ -4,12 +4,12 @@
 # Copyright (c) 2023, Tri Dao, Albert Gu.
 
 import math
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange, repeat
-
-from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn
+from mamba_ssm.ops.selective_scan_interface import mamba_inner_fn, selective_scan_fn
 
 try:
     from causal_conv1d import causal_conv1d_fn, causal_conv1d_update

samba_pytorch/modules/multiscale_retention.py

@@ -12,9 +12,6 @@
 import torch
 import torch.nn as nn
 from einops import rearrange
-from transformers.activations import ACT2FN
-from transformers.cache_utils import Cache
-
 from fla.modules import FusedRMSNormSwishGate, RMSNorm, ShortConvolution
 from fla.modules.rotary import RotaryEmbedding
 from fla.ops.retention import (
@@ -23,10 +20,11 @@
     fused_recurrent_retention,
     parallel_retention,
 )
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache
 
 
 class MultiScaleRetention(nn.Module):
-
     def __init__(
         self,
         mode: str = "fused_chunk",

samba_pytorch/modules/rmsnorm.py

@@ -4,7 +4,6 @@
 # Copyright (c) 2022, Tri Dao.
 # Adapted from https://github.com/NVIDIA/apex/blob/master/apex/contrib/layer_norm/layer_norm.py AND https://github.com/Dao-AILab/flash-attention/blob/7a983df74215e035e566e37125b0a71e3618f39d/flash_attn/ops/layer_norm.py#L16
 
-import torch
 import dropout_layer_norm
 import torch
 from torch.nn import init

samba_pytorch/modules/rotary.py

@@ -6,7 +6,6 @@
 from typing import Optional, Union
 
 import torch
-
 import triton
 import triton.language as tl
 

samba_pytorch/samba.py

@@ -5,29 +5,28 @@
 # see LICENSE file at https://github.com/Lightning-AI/litgpt/blob/main/LICENSE
 
 import math
+from functools import partial
 from typing import Any, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
-from typing_extensions import Self
 from lit_gpt.config import Config
+from torch import Tensor
+from typing_extensions import Self
 from xformers.ops import SwiGLU
+
 from .fused_rotary_embedding import apply_rotary_emb_func
-from torch import Tensor
 from .mamba_simple import Mamba
-from functools import partial
 
 try:
     from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
 except ImportError:
     RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
-from .gla import GatedLinearAttention
-from .multiscale_retention import MultiScaleRetention
-from einops import rearrange
-import torch.nn.functional as F
-
 from causal_conv1d import causal_conv1d_fn
+from einops import rearrange
 
+from .gla import GatedLinearAttention
+from .multiscale_retention import MultiScaleRetention
 
 RoPECache = Tuple[torch.Tensor, torch.Tensor]
 KVCache = Tuple[torch.Tensor, torch.Tensor]
@@ -396,7 +395,6 @@ def forward(
         input_pos: Optional[torch.Tensor] = None,
         kv_cache: Optional[KVCache] = None,
     ) -> Tuple[torch.Tensor, Optional[KVCache]]:
-
         n_1 = self.norm_1(x)
 
         if self.use_mamba: