muscriptor 0.1.0__py3-none-any.whl

muscriptor/modules/conditioners.py

@@ -0,0 +1,321 @@
+"""Minimal conditioners for muscriptor inference.
+
+Contains only the classes needed to run the transcription model:
+- ConditioningAttributes, ConditionType, WavCondition
+- MelSpectrogramConditioner (audio → mel → linear projection)
+- ClassConditioner (class index → embedding)
+- ConditioningProvider
+- nullify_all_conditions (for CFG at inference)
+"""
+
+import time
+from collections import defaultdict
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import NamedTuple, Any
+import torch
+from torch import nn
+from torch.nn import functional as F
+from einops import rearrange
+
+from muscriptor.modules.mel_spectrogram import _MelSpectrogram
+from muscriptor.utils.sampling import length_to_mask
+
+
+ConditionType = tuple[torch.Tensor, torch.Tensor]  # (embedding [B, T, D], mask [B, T])
+
+
+class WavCondition(NamedTuple):
+    wav: torch.Tensor
+    length: torch.Tensor
+    sample_rate: list[int]
+    path: list[str | None] = []
+    seek_time: list[float | None] = []
+
+
+@dataclass
+class ConditioningAttributes:
+    text: dict[str, str | None] = field(default_factory=dict)
+    wav: dict[str, WavCondition] = field(default_factory=dict)
+    joint_embed: dict[str, Any] = field(default_factory=dict)
+    symbolic: dict[str, Any] = field(default_factory=dict)
+
+    def __getitem__(self, item):
+        return getattr(self, item)
+
+    @property
+    def text_attributes(self):
+        return self.text.keys()
+
+    @property
+    def wav_attributes(self):
+        return self.wav.keys()
+
+    @property
+    def joint_embed_attributes(self):
+        return self.joint_embed.keys()
+
+    @property
+    def symbolic_attributes(self):
+        return self.symbolic.keys()
+
+    @property
+    def attributes(self):
+        return {
+            "text": self.text_attributes,
+            "wav": self.wav_attributes,
+            "joint_embed": self.joint_embed_attributes,
+            "symbolic": self.symbolic_attributes,
+        }
+
+    @classmethod
+    def condition_types(cls) -> list:
+        return ["text", "wav", "joint_embed", "symbolic"]
+
+
+def nullify_wav(cond: WavCondition) -> WavCondition:
+    B = cond.wav.shape[0]
+    return WavCondition(
+        wav=torch.zeros(*cond.wav.shape[:-1], 1, device=cond.wav.device),
+        length=torch.zeros(B, dtype=cond.length.dtype, device=cond.wav.device),
+        sample_rate=cond.sample_rate,
+        path=[None] * B,
+        seek_time=[None] * B,
+    )
+
+
+def nullify_all_conditions(
+    samples: list[ConditioningAttributes],
+) -> list[ConditioningAttributes]:
+    """Return a copy of ``samples`` with every wav/text condition nulled out.
+
+    Used to build the unconditional batch for classifier-free guidance.
+    """
+    samples = deepcopy(samples)
+    for sample in samples:
+        for k in list(sample.wav):
+            sample.wav[k] = nullify_wav(sample.wav[k])
+        for k in list(sample.text):
+            sample.text[k] = None
+    return samples
+
+
+class MelSpectrogramConditioner(nn.Module):
+    """Log mel spectrogram conditioner. Projects mel bins to transformer dim."""
+
+    def __init__(
+        self,
+        output_dim: int,
+        device: torch.device | str,
+        sample_rate: int,
+        n_fft: int = 2048,
+        frame_rate: int = 100,
+        n_mel_bins: int = 512,
+        normalize_audio: bool = False,
+        log_scale: bool = True,
+        eps: float = 1e-6,
+        # unused arg kept for config compatibility
+        fine_frame_rate: int = None,
+    ):
+        self.fine_frame_rate_ratio = 1
+        if fine_frame_rate is not None:
+            assert fine_frame_rate % frame_rate == 0
+            self.fine_frame_rate_ratio = fine_frame_rate // frame_rate
+
+        super().__init__()
+        self.dim = n_mel_bins * self.fine_frame_rate_ratio
+        self.output_dim = output_dim
+        self.output_proj = nn.Linear(self.dim, output_dim)
+        self.device = device
+        self.sample_rate = sample_rate
+        self.frame_rate = frame_rate
+        self.normalize_audio = normalize_audio
+        self.log_scale = log_scale
+        self.eps = eps
+
+        if self.fine_frame_rate_ratio == 1:
+            assert sample_rate % frame_rate == 0
+            self.hop_length = sample_rate // frame_rate
+        else:
+            assert sample_rate % fine_frame_rate == 0
+            self.hop_length = sample_rate // fine_frame_rate
+
+        self.mel_spec_transform = _MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=n_fft,
+            hop_length=self.hop_length,
+            n_mels=n_mel_bins,
+            power=1.0,
+            center=True,
+            pad_mode="reflect",
+        ).to(device)
+
+    def tokenize(self, x: WavCondition) -> WavCondition:
+        wav, length, sample_rate, path, seek_time = x
+        assert length is not None
+        return WavCondition(
+            wav.to(self.device), length.to(self.device), sample_rate, path, seek_time
+        )
+
+    def _mel_embedding(self, x: WavCondition) -> torch.Tensor:
+        if x.wav.shape[-1] == 1:
+            return torch.zeros(x.wav.shape[0], 1, self.dim, device=self.device)
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            wav = x.wav
+            if self.normalize_audio:
+                wav = wav / (wav.abs().max(dim=-1, keepdim=True).values + 1e-8)
+            mel = self.mel_spec_transform(wav)
+            mel = rearrange(mel, "b 1 d t -> b t d")
+            if self.fine_frame_rate_ratio > 1:
+                mel = rearrange(
+                    mel[:, :-1], "b (t f) d -> b t (f d)", f=self.fine_frame_rate_ratio
+                )
+            if self.log_scale:
+                mel = torch.log(mel + self.eps)
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        print(
+            f"[muscriptor] mel-spec ({wav.shape[0]} × {wav.shape[-1]} samples): "
+            f"{time.perf_counter() - t0:.3f}s"
+        )
+        return mel
+
+    def forward(self, x: WavCondition) -> ConditionType:
+        _, lengths, *_ = x
+        with torch.no_grad():
+            embeds = self._mel_embedding(x)
+        embeds = embeds.to(self.output_proj.weight)
+        embeds = self.output_proj(embeds)
+
+        if lengths is not None:
+            lengths = lengths / (self.sample_rate // self.frame_rate)
+            mask = length_to_mask(lengths, max_len=embeds.shape[1]).int()
+        else:
+            mask = torch.ones_like(embeds[..., 0])
+        mask_f = mask.float().unsqueeze(-1).to(embeds.device)
+        embeds = embeds * mask_f
+        return embeds, mask
+
+
+class ClassConditioner(nn.Module):
+    """Conditioner that embeds class indices (e.g., instrument group, dataset name)."""
+
+    def __init__(
+        self,
+        num_classes: int,
+        output_dim: int,
+        device: torch.device | str = "cpu",
+    ):
+        super().__init__()
+        self.device = device
+        self.embed = nn.Embedding(num_classes + 1, output_dim).to(device)
+        self.pad_idx = 0
+
+    def tokenize(self, x: list[str | None]) -> torch.Tensor:
+        int_x = [list(map(int, s.split())) if s is not None else [-1] for s in x]
+        max_len = max(len(xi) for xi in int_x)
+        int_x = [xi + [-1] * (max_len - len(xi)) for xi in int_x]
+        int_x = 1 + torch.LongTensor(int_x).to(self.device)
+        return int_x
+
+    def forward(self, inputs: torch.Tensor) -> ConditionType:
+        embeds = self.embed(inputs + 1)
+        mask = torch.ones_like(embeds[..., 0])
+        return embeds, mask
+
+
+def collate_wavs(
+    samples: list[ConditioningAttributes], wav_conditions: list[str]
+) -> dict[str, WavCondition]:
+    """Collate wav conditions from a list of ConditioningAttributes."""
+    wavs = defaultdict(list)
+    lengths = defaultdict(list)
+    sample_rates: dict[str, list] = defaultdict(list)
+    paths: dict[str, list] = defaultdict(list)
+    seek_times: dict[str, list] = defaultdict(list)
+    out: dict[str, WavCondition] = {}
+
+    for sample in samples:
+        for attribute in wav_conditions:
+            wav, length, sample_rate, path, seek_time = sample.wav[attribute]
+            assert wav.dim() == 3
+            B, K, T = wav.shape
+            assert B == 1
+            if K == 2:
+                wav = wav.mean(1, keepdim=True)
+            wavs[attribute].append(wav)
+            lengths[attribute].append(length)
+            sample_rates[attribute].extend(sample_rate)
+            paths[attribute].extend(path)
+            seek_times[attribute].extend(seek_time)
+
+    for attribute in wav_conditions:
+        # Stack along batch dim
+        all_wavs = wavs[attribute]
+        max_len = max(w.shape[-1] for w in all_wavs)
+        padded = torch.cat(
+            [F.pad(w, (0, max_len - w.shape[-1])) for w in all_wavs], dim=0
+        )
+        out[attribute] = WavCondition(
+            padded,
+            torch.cat(lengths[attribute]),
+            sample_rates[attribute],
+            paths[attribute],
+            seek_times[attribute],
+        )
+    return out
+
+
+class ConditioningProvider(nn.Module):
+    """Runs all conditioners and returns a dict of condition tensors."""
+
+    def __init__(
+        self,
+        conditioners: dict[str, nn.Module],
+        device: torch.device | str = "cpu",
+    ):
+        super().__init__()
+        self.device = device
+        self.conditioners = nn.ModuleDict(conditioners)
+
+    @property
+    def text_conditions(self):
+        return [
+            k for k, v in self.conditioners.items() if isinstance(v, ClassConditioner)
+        ]
+
+    @property
+    def wav_conditions(self):
+        return [
+            k
+            for k, v in self.conditioners.items()
+            if isinstance(v, MelSpectrogramConditioner)
+        ]
+
+    def tokenize(self, inputs: list[ConditioningAttributes]) -> dict[str, Any]:
+        output = {}
+        # Collate text conditions
+        text_batch: dict[str, list[str | None]] = defaultdict(list)
+        for sample in inputs:
+            for cond in self.text_conditions:
+                text_batch[cond].append(sample.text.get(cond))
+        for attr, batch in text_batch.items():
+            output[attr] = self.conditioners[attr].tokenize(batch)
+
+        # Collate wav conditions
+        if self.wav_conditions:
+            wav_batch = collate_wavs(inputs, self.wav_conditions)
+            for attr, wav_cond in wav_batch.items():
+                output[attr] = self.conditioners[attr].tokenize(wav_cond)
+
+        return output
+
+    def forward(self, tokenized: dict[str, Any]) -> dict[str, ConditionType]:
+        output = {}
+        for attribute, inputs in tokenized.items():
+            condition, mask = self.conditioners[attribute](inputs)
+            output[attribute] = (condition, mask)
+        return output

muscriptor/modules/mel_spectrogram.py

@@ -0,0 +1,106 @@
+"""Standalone pure-torch equivalent of torchaudio.transforms.MelSpectrogram.
+
+Matches torchaudio with mel_scale='htk', norm=None, win_length == n_fft.
+"""
+
+import torch
+from torch import nn
+
+
+def _hz_to_mel_htk(freq: torch.Tensor) -> torch.Tensor:
+    return 2595.0 * torch.log10(1.0 + freq / 700.0)
+
+
+def _mel_to_hz_htk(mel: torch.Tensor) -> torch.Tensor:
+    return 700.0 * (10 ** (mel / 2595.0) - 1.0)
+
+
+def melscale_fbanks(
+    n_freqs: int,
+    f_min: float,
+    f_max: float,
+    n_mels: int,
+    sample_rate: int,
+) -> torch.Tensor:
+    """Triangular mel filterbank matching torchaudio.functional.melscale_fbanks
+    with mel_scale='htk' and norm=None. Returns a tensor of shape [n_freqs, n_mels]."""
+    all_freqs = torch.linspace(0, sample_rate // 2, n_freqs)
+    m_min = _hz_to_mel_htk(torch.tensor(float(f_min)))
+    m_max = _hz_to_mel_htk(torch.tensor(float(f_max)))
+    m_pts = torch.linspace(m_min.item(), m_max.item(), n_mels + 2)
+    f_pts = _mel_to_hz_htk(m_pts)
+
+    f_diff = f_pts[1:] - f_pts[:-1]
+    slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1)
+    down_slopes = -slopes[:, :-2] / f_diff[:-1]
+    up_slopes = slopes[:, 2:] / f_diff[1:]
+    return torch.maximum(torch.zeros(()), torch.minimum(down_slopes, up_slopes))
+
+
+class _Spectrogram(nn.Module):
+    """Holds the STFT window so the safetensors key
+    `...mel_spec_transform.spectrogram.window` round-trips."""
+
+    def __init__(self, n_fft: int):
+        super().__init__()
+        self.register_buffer("window", torch.hann_window(n_fft))
+
+
+class _MelScale(nn.Module):
+    """Holds the mel filterbank so the safetensors key
+    `...mel_spec_transform.mel_scale.fb` round-trips."""
+
+    def __init__(self, fb: torch.Tensor):
+        super().__init__()
+        self.register_buffer("fb", fb)
+
+
+class _MelSpectrogram(nn.Module):
+    """Pure-torch equivalent of torchaudio.transforms.MelSpectrogram
+    (htk mel scale, no Slaney norm, win_length == n_fft)."""
+
+    def __init__(
+        self,
+        sample_rate: int,
+        n_fft: int,
+        hop_length: int,
+        n_mels: int,
+        power: float = 2.0,
+        center: bool = True,
+        pad_mode: str = "reflect",
+    ):
+        super().__init__()
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.power = power
+        self.center = center
+        self.pad_mode = pad_mode
+
+        self.spectrogram = _Spectrogram(n_fft)
+        fb = melscale_fbanks(
+            n_freqs=n_fft // 2 + 1,
+            f_min=0.0,
+            f_max=sample_rate / 2.0,
+            n_mels=n_mels,
+            sample_rate=sample_rate,
+        )
+        self.mel_scale = _MelScale(fb)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        leading = x.shape[:-1]
+        x = x.reshape(-1, x.shape[-1])
+        spec = torch.stft(
+            x,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.n_fft,
+            window=self.spectrogram.window,
+            center=self.center,
+            pad_mode=self.pad_mode,
+            return_complex=True,
+            normalized=False,
+            onesided=True,
+        )
+        spec = spec.abs() ** self.power
+        mel = torch.matmul(spec.transpose(-1, -2), self.mel_scale.fb).transpose(-1, -2)
+        return mel.reshape(*leading, *mel.shape[-2:])

muscriptor/modules/streaming.py

@@ -0,0 +1,68 @@
+"""Stateful module API.
+
+Each :class:`StatefulModule` exposes :meth:`init_state` returning a dict of
+per-module tensors. :func:`init_states` walks an ``nn.Module`` tree, calls
+``init_state`` on every stateful submodule, and returns a ``dict[name -> state]``
+that callers thread through ``forward`` via a ``model_state`` argument.
+
+State is mutated only by :meth:`increment_step` (called explicitly via
+:func:`increment_steps`) and by ``forward`` writing into preallocated buffers
+at known offsets.  No magic context manager, no implicit per-module storage.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+from torch import nn
+
+
+State = dict[str, Any]
+ModelState = dict[str, State]
+
+
+class StatefulModule(ABC, nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._module_absolute_name: str | None = None
+
+    @abstractmethod
+    def init_state(self, batch_size: int, sequence_length: int) -> State:
+        raise NotImplementedError
+
+    def increment_step(self, state: State, increment: int = 1) -> None:
+        pass
+
+    def get_state(self, model_state: ModelState | None) -> State | None:
+        if model_state is None or self._module_absolute_name is None:
+            return None
+        return model_state.get(self._module_absolute_name)
+
+
+def init_states(model: nn.Module, batch_size: int, sequence_length: int) -> ModelState:
+    """Allocate state for every :class:`StatefulModule` reachable from ``model``.
+
+    Side effect: each stateful submodule has its ``_module_absolute_name`` set
+    so subsequent ``get_state`` calls can find its slot.
+    """
+    result: ModelState = {}
+    for module_name, module in model.named_modules():
+        if isinstance(module, StatefulModule):
+            module._module_absolute_name = module_name
+            result[module_name] = module.init_state(batch_size, sequence_length)
+    return result
+
+
+def increment_steps(
+    model: nn.Module, model_state: ModelState, increment: int = 1
+) -> None:
+    """Bump the step counter for every stateful submodule of ``model``.
+
+    Uses each module's ``_module_absolute_name`` (set by :func:`init_states`)
+    to look up its slot, so this works on subtrees even when ``init_states``
+    was called on a different root.
+    """
+    for _, module in model.named_modules():
+        if (
+            isinstance(module, StatefulModule)
+            and module._module_absolute_name is not None
+        ):
+            module.increment_step(model_state[module._module_absolute_name], increment)

muscriptor/modules/transformer.py

@@ -0,0 +1,202 @@
+"""Causal streaming transformer for muscriptor inference."""
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from muscriptor.modules.streaming import ModelState, State, StatefulModule
+
+
+def create_sin_embedding(
+    positions: torch.Tensor,
+    dim: int,
+    max_period: float = 10000,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    assert dim % 2 == 0
+    half_dim = dim // 2
+    positions = positions.to(dtype)
+    adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1)
+    max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype)
+    phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))
+    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+class StreamingMultiheadAttention(StatefulModule):
+    """Causal multi-head self-attention with a preallocated KV cache."""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dim_per_head = embed_dim // num_heads
+
+        in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False, **factory_kwargs)
+        self.in_proj_weight = in_proj.weight
+        self.in_proj_bias = in_proj.bias
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False, **factory_kwargs)
+
+    def init_state(self, batch_size: int, sequence_length: int) -> State:
+        weight = self.in_proj_weight
+        return {
+            "cache": torch.full(
+                (2, batch_size, sequence_length, self.num_heads, self.dim_per_head),
+                float("nan"),
+                device=weight.device,
+                dtype=weight.dtype,
+            ),
+            "offset": torch.zeros(1, dtype=torch.long, device=weight.device),
+        }
+
+    def increment_step(self, state: State, increment: int = 1) -> None:
+        state["offset"] = state["offset"] + increment
+
+    def _complete_kv(self, k, v, state: State | None):
+        if state is None:
+            return k, v
+        cache = state["cache"]
+        end = int(state["offset"].item())
+        T = k.shape[1]
+        cache[0, :, end : end + T] = k
+        cache[1, :, end : end + T] = v
+        return cache[0, :, : end + T], cache[1, :, : end + T]
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        model_state: ModelState | None = None,
+    ):
+        state = self.get_state(model_state)
+        projected = nn.functional.linear(query, self.in_proj_weight)
+        packed = rearrange(projected, "b t (p h d) -> b t p h d", p=3, h=self.num_heads)
+        q, k, v = packed.unbind(dim=2)
+
+        k, v = self._complete_kv(k, v, state)
+        dtype = q.dtype
+
+        q_t = q.transpose(1, 2)
+        k_t = k.transpose(1, 2)
+        v_t = v.transpose(1, 2)
+
+        # Explicit bottom-right causal mask so streaming decode steps
+        # (T_q=1, T_k=cache_len) attend to all past tokens. PyTorch's
+        # is_causal=True uses top-left alignment and would mask out all
+        # cached tokens except position 0 when T_q < T_k.
+        T_q, T_k = q_t.shape[2], k_t.shape[2]
+        mask = torch.ones(T_q, T_k, dtype=torch.bool, device=query.device).tril(
+            T_k - T_q
+        )
+        attn_bias = torch.zeros(T_q, T_k, dtype=q.dtype, device=query.device)
+        attn_bias.masked_fill_(~mask, float("-inf"))
+        x = F.scaled_dot_product_attention(
+            q_t, k_t, v_t, attn_mask=attn_bias, dropout_p=0.0
+        )
+        x = x.transpose(1, 2).to(dtype)
+
+        x = rearrange(x, "b t h d -> b t (h d)")
+        x = self.out_proj(x)
+        return x
+
+
+class StreamingTransformerLayer(nn.Module):
+    """Pre-norm transformer block: self-attention + GELU FFN."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_feedforward: int = 2048,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.self_attn = StreamingMultiheadAttention(
+            embed_dim=d_model, num_heads=num_heads, **factory_kwargs
+        )
+        self.norm1 = nn.LayerNorm(d_model, eps=1e-5, **factory_kwargs)
+        self.norm2 = nn.LayerNorm(d_model, eps=1e-5, **factory_kwargs)
+        self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False, **factory_kwargs)
+        self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False, **factory_kwargs)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        model_state: ModelState | None = None,
+    ):
+        x = x + self.self_attn(self.norm1(x), model_state=model_state)
+        x = x + self.linear2(F.gelu(self.linear1(self.norm2(x))))
+        return x
+
+
+class StreamingTransformer(StatefulModule):
+    """Stack of causal streaming transformer layers with sinusoidal positions."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_layers: int,
+        dim_feedforward: int = 2048,
+        max_period: float = 10_000,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        assert d_model % num_heads == 0
+        self.max_period = max_period
+        self.layers = nn.ModuleList(
+            [
+                StreamingTransformerLayer(
+                    d_model=d_model,
+                    num_heads=num_heads,
+                    dim_feedforward=dim_feedforward,
+                    device=device,
+                    dtype=dtype,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def init_state(self, batch_size: int, sequence_length: int) -> State:
+        device = self.layers[0].norm2.weight.device
+        return {
+            "offsets": torch.zeros(batch_size, dtype=torch.long, device=device),
+        }
+
+    def increment_step(self, state: State, increment: int = 1) -> None:
+        state["offsets"] = state["offsets"] + increment
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        prepend_length: int = 0,
+        model_state: ModelState | None = None,
+    ):
+        del prepend_length  # unused; positions come from state['offsets']
+        B, T, C = x.shape
+        state = self.get_state(model_state)
+        offsets = (
+            state["offsets"]
+            if state is not None
+            else torch.zeros(B, dtype=torch.long, device=x.device)
+        )
+
+        positions = torch.arange(T, device=x.device).view(1, -1, 1)
+        positions = positions + offsets.view(-1, 1, 1)
+        pos_emb = create_sin_embedding(
+            positions, C, max_period=self.max_period, dtype=x.dtype
+        )
+        x = x + pos_emb * (positions >= 0).float()
+
+        for layer in self.layers:
+            x = layer(x, model_state=model_state)
+        return x