Models

def __init__(self, in_features: int, out_features: int):
    super().__init__(
        in_channels=1,
        out_channels=1,
        kernel_size=in_features + out_features - 1,
        padding=out_features - 1,
        bias=False,
    )

def __init__(
    self,
    *,
    n_chan_input: int = 1,
    n_chan_layers: tuple[int, ...] = (40, 30, 30, 10, 3),
    n_prefilt_layers: int = 3,
    prefilt_kernel_size: int = 39,
    residual: bool = True,
    n_bins_in: int = 219,
    output_dim: int = 384,
    activation_fn: Literal["relu", "silu", "leaky"] = "leaky",
    a_lrelu: float = 0.3,
    p_dropout: float = 0.2,
):
    super().__init__()

    activation_layer: Callable[[], nn.Module]
    if activation_fn == "relu":
        activation_layer = nn.ReLU
    elif activation_fn == "silu":
        activation_layer = nn.SiLU
    elif activation_fn == "leaky":
        activation_layer = partial(nn.LeakyReLU, negative_slope=a_lrelu)
    else:
        raise ValueError(f"unsupported activation_fn={activation_fn!r}")

    n_ch = list(n_chan_layers)
    if len(n_ch) < 5:
        n_ch.append(1)

    self.layernorm = nn.LayerNorm(normalized_shape=[n_chan_input, n_bins_in])

    prefilt_padding = prefilt_kernel_size // 2
    self.conv1 = nn.Sequential(
        nn.Conv1d(
            in_channels=n_chan_input,
            out_channels=n_ch[0],
            kernel_size=prefilt_kernel_size,
            padding=prefilt_padding,
            stride=1,
        ),
        activation_layer(),
        nn.Dropout(p=p_dropout),
    )
    self.n_prefilt_layers = n_prefilt_layers
    self.prefilt_layers = nn.ModuleList(
        [
            nn.Sequential(
                nn.Conv1d(
                    in_channels=n_ch[0],
                    out_channels=n_ch[0],
                    kernel_size=prefilt_kernel_size,
                    padding=prefilt_padding,
                    stride=1,
                ),
                activation_layer(),
                nn.Dropout(p=p_dropout),
            )
            for _ in range(n_prefilt_layers - 1)
        ]
    )
    self.residual = residual

    conv_layers: list[nn.Module] = []
    for i in range(len(n_ch) - 1):
        conv_layers.extend(
            [
                nn.Conv1d(
                    in_channels=n_ch[i],
                    out_channels=n_ch[i + 1],
                    kernel_size=1,
                    padding=0,
                    stride=1,
                ),
                activation_layer(),
                nn.Dropout(p=p_dropout),
            ]
        )
    self.conv_layers = nn.Sequential(*conv_layers)

    self.flatten = nn.Flatten(start_dim=1)
    self.fc = ToeplitzLinear(n_bins_in * n_ch[-1], output_dim)
    self.final_norm = nn.Softmax(dim=-1)

def __init__(self) -> None:
    super().__init__()
    self.conv = nn.Conv1d(1, 1, 39, stride=3)
    self.linear = nn.Linear(72, 1)

def __init__(self, cfg: PestoParams):
    super().__init__()
    self.cfg = cfg
    self.encoder = Resnet1d(
        n_chan_input=cfg.n_chan_input,
        n_chan_layers=cfg.n_chan_layers,
        n_prefilt_layers=cfg.n_prefilt_layers,
        prefilt_kernel_size=cfg.prefilt_kernel_size,
        residual=cfg.residual,
        n_bins_in=cfg.n_bins_in,
        output_dim=cfg.output_dim,
        activation_fn=cfg.activation_fn,
        a_lrelu=cfg.a_lrelu,
        p_dropout=cfg.p_dropout,
    )
    self.confidence = ConfidenceClassifier()

    self.register_buffer("shift", torch.zeros((), dtype=torch.float), persistent=True)

def __init__(
    self,
    dim: int,
    dim_head: int,
    n_head: int,
    rotary_embed_f: RotaryEmbedding,
    rotary_embed_t: RotaryEmbedding,
    dropout: float,
):
    super().__init__()
    self.attnF = Attention(
        dim, heads=n_head, dim_head=dim_head, dropout=dropout, rotary_embed=rotary_embed_f
    )
    self.ffF = FeedForward(dim, dropout=dropout)
    self.attnT = Attention(
        dim, heads=n_head, dim_head=dim_head, dropout=dropout, rotary_embed=rotary_embed_t
    )
    self.ffT = FeedForward(dim, dropout=dropout)

def __init__(self, input_dim: int):
    super().__init__()
    self.beat_downbeat_lin = nn.Linear(input_dim, 2)

def __init__(self, input_dim: int):
    super().__init__()
    self.beat_downbeat_lin = nn.Linear(input_dim, 2)

def __init__(self, cfg: BeatThisParams):
    super().__init__()
    self.spect_dim = cfg.spect_dim

    max_frames = cfg.chunk_size // cfg.log_mel_hop_length
    rotary_embed_t = RotaryEmbedding(
        seq_len=max_frames,  # by default 1500 frames * 441 = 661500 samples
        dim_head=cfg.head_dim,
        dtype=cfg.rotary_embed_dtype,
    )

    # NOTE: removed rearrange from original impl to standardise input as (B, F, T)
    stem = nn.Sequential(
        OrderedDict(
            bn1d=nn.BatchNorm1d(cfg.spect_dim),
            add_channel=Rearrange("b f t -> b 1 f t"),
            conv2d=nn.Conv2d(
                in_channels=1,
                out_channels=cfg.stem_dim,
                kernel_size=(4, 3),
                stride=(4, 1),
                padding=(0, 1),
                bias=False,
            ),
            bn2d=nn.BatchNorm2d(cfg.stem_dim),
            activation=nn.GELU(),
        )
    )

    spect_dim = cfg.spect_dim // 4
    dim = cfg.stem_dim
    frontend_blocks = []
    for _ in range(3):
        rotary_embed_f = RotaryEmbedding(
            seq_len=spect_dim,
            dim_head=cfg.head_dim,
            dtype=cfg.rotary_embed_dtype,
        )
        frontend_blocks.append(
            nn.Sequential(
                OrderedDict(
                    partial=(
                        PartialFTTransformer(
                            dim=dim,
                            dim_head=cfg.head_dim,
                            n_head=dim // cfg.head_dim,
                            rotary_embed_f=rotary_embed_f,
                            rotary_embed_t=rotary_embed_t,
                            dropout=cfg.dropout_frontend,
                        )
                        if cfg.partial_transformers
                        else nn.Identity()
                    ),
                    conv2d=nn.Conv2d(
                        in_channels=dim,
                        out_channels=dim * 2,
                        kernel_size=(2, 3),
                        stride=(2, 1),
                        padding=(0, 1),
                        bias=False,
                    ),
                    norm=nn.BatchNorm2d(dim * 2),
                    activation=nn.GELU(),
                )
            )
        )
        dim *= 2
        spect_dim //= 2

    self.frontend = nn.Sequential(
        OrderedDict(
            stem=stem,
            blocks=nn.Sequential(*frontend_blocks),
            concat=Rearrange("b c f t -> b t (c f)"),
            linear=nn.Linear(dim * spect_dim, cfg.transformer_dim),
        )
    )

    n_heads = cfg.transformer_dim // cfg.head_dim
    # TODO check if this is really equivalent
    self.transformer_blocks = Transformer(
        dim=cfg.transformer_dim,
        depth=cfg.n_layers,
        dim_head=cfg.head_dim,
        heads=n_heads,
        attn_dropout=cfg.dropout_transformer,
        ff_dropout=cfg.dropout_transformer,
        ff_mult=cfg.ff_mult,
        norm_output=True,
        rotary_embed=rotary_embed_t,
        transformer_residual_dtype=cfg.transformer_residual_dtype,
    )

    if cfg.sum_head:
        self.task_heads = SumHead(cfg.transformer_dim)
    else:
        self.task_heads = Head(cfg.transformer_dim)