autogluon.timeseries 1.4.1b20250924__tar.gz → 1.4.1b20251011__tar.gz

{autogluon.timeseries-1.4.1b20250924 → autogluon.timeseries-1.4.1b20251011}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: autogluon.timeseries
-Version: 1.4.1b20250924
+Version: 1.4.1b20251011
 Summary: Fast and Accurate ML in 3 Lines of Code
 Home-page: https://github.com/autogluon/autogluon
 Author: AutoGluon Community
@@ -35,6 +35,7 @@ Classifier: Topic :: Scientific/Engineering :: Image Recognition
 Requires-Python: >=3.9, <3.13
 Description-Content-Type: text/markdown
 Provides-Extra: tests
+Provides-Extra: toto
 Provides-Extra: all
 License-File: ../LICENSE
 License-File: ../NOTICE

{autogluon.timeseries-1.4.1b20250924 → autogluon.timeseries-1.4.1b20251011}/setup.py

@@ -54,10 +54,14 @@ extras_require = {
         "flaky>=3.7,<4",
         "pytest-timeout>=2.1,<3",
     ],
+    "toto": [
+        "einops>=0.7,<1",
+        "rotary-embedding-torch>=0.8,<1",
+    ],
 }
 
 # chronos-openvino and chronos-onnx are deprecated, and will be removed in a future version
-extras_require["all"] = []
+extras_require["all"] = list(set.union(*(set(extras_require[extra]) for extra in ["toto"])))
 install_requires = ag.get_dependency_version_ranges(install_requires)
 
 if __name__ == "__main__":

{autogluon.timeseries-1.4.1b20250924 → autogluon.timeseries-1.4.1b20251011}/src/autogluon/timeseries/models/__init__.py

@@ -28,6 +28,7 @@ from .local import (
     ZeroModel,
 )
 from .registry import ModelRegistry
+from .toto import TotoModel
 
 __all__ = [
     "ADIDAModel",
@@ -56,6 +57,7 @@ __all__ = [
     "TemporalFusionTransformerModel",
     "ThetaModel",
     "TiDEModel",
+    "TotoModel",
     "WaveNetModel",
     "ZeroModel",
 ]

autogluon.timeseries-1.4.1b20251011/src/autogluon/timeseries/models/toto/__init__.py

@@ -0,0 +1,3 @@
+from .model import TotoModel
+
+__all__ = ["TotoModel"]

autogluon.timeseries-1.4.1b20251011/src/autogluon/timeseries/models/toto/_internal/__init__.py

@@ -0,0 +1,9 @@
+from .backbone import TotoBackbone
+from .dataset import MaskedTimeseries
+from .forecaster import TotoForecaster
+
+__all__ = [
+    "MaskedTimeseries",
+    "TotoBackbone",
+    "TotoForecaster",
+]

autogluon.timeseries-1.4.1b20251011/src/autogluon/timeseries/models/toto/_internal/backbone/__init__.py

@@ -0,0 +1,3 @@
+from .backbone import TotoBackbone
+
+__all__ = ["TotoBackbone"]

autogluon.timeseries-1.4.1b20251011/src/autogluon/timeseries/models/toto/_internal/backbone/attention.py

@@ -0,0 +1,197 @@
+# Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+#
+# This product includes software developed at Datadog (https://www.datadoghq.com/)
+# Copyright 2025 Datadog, Inc.
+
+import logging
+from enum import Enum
+from typing import Optional, Union
+
+import torch
+from einops import rearrange
+from torch.nn.functional import scaled_dot_product_attention
+
+from .rope import TimeAwareRotaryEmbedding
+
+log = logging.getLogger(__name__)
+
+
+class AttentionAxis(Enum):
+    TIME = 1
+    SPACE = 2
+
+
+class BaseMultiheadAttention(torch.nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float,
+        rotary_emb: Optional[TimeAwareRotaryEmbedding],
+        use_memory_efficient_attention: bool,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        assert embed_dim % num_heads == 0, "Embedding dimension must be divisible by number of heads."
+        self.head_dim = embed_dim // num_heads
+        self.rotary_emb = rotary_emb
+
+        # We allocate a single tensor for the q, k, and v projection matrices,
+        # multiply them with the inputs, and then split the projected tensors into q, k, and v using unbind.
+        # This reduces overhead a bit vs. having multiple separate Linear layers,
+        # which need to be initialized, tracked by the optimizer, etc.
+        self.wQKV = torch.nn.Linear(embed_dim, embed_dim * 3)
+        self.dropout = dropout
+        self.use_memory_efficient_attention = use_memory_efficient_attention
+        self.wO = torch.nn.Linear(embed_dim, embed_dim)
+
+        assert not self.use_memory_efficient_attention, (
+            "xformers is not available, so use_memory_efficient_attention must be False"
+        )
+
+        if not hasattr(self, "attention_axis") or self.attention_axis not in (AttentionAxis.TIME, AttentionAxis.SPACE):
+            raise ValueError("Child class must define attention_axis as AttentionAxis.TIME or AttentionAxis.SPACE.")
+
+    def rearrange_inputs(self, inputs: torch.Tensor) -> torch.Tensor:
+        pattern = (
+            "batch variate seq_len embed_dim -> (batch variate) seq_len embed_dim"
+            if self.attention_axis == AttentionAxis.TIME
+            else "batch variate seq_len embed_dim -> (batch seq_len) variate embed_dim"
+        )
+
+        return rearrange(inputs, pattern)
+
+    def get_qkv(
+        self,
+        inputs: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        pattern: str = ""
+        if self.attention_axis == AttentionAxis.TIME and self.use_memory_efficient_attention:
+            pattern = "batch_X_variate seq_len (qkv head_dim n_heads) -> qkv batch_X_variate seq_len n_heads head_dim"
+        elif self.attention_axis == AttentionAxis.TIME and not self.use_memory_efficient_attention:
+            pattern = "batch_X_variate seq_len (qkv head_dim n_heads) -> qkv batch_X_variate n_heads seq_len head_dim"
+        elif self.attention_axis == AttentionAxis.SPACE and self.use_memory_efficient_attention:
+            pattern = "batch_X_seq_len variate (qkv head_dim n_heads) -> qkv batch_X_seq_len variate n_heads head_dim"
+        elif self.attention_axis == AttentionAxis.SPACE and not self.use_memory_efficient_attention:
+            pattern = "batch_X_seq_len variate (qkv head_dim n_heads) -> qkv batch_X_seq_len n_heads variate head_dim"
+
+        assert pattern
+        qkv = self.wQKV(inputs.contiguous())
+        return rearrange(qkv, pattern, qkv=3, head_dim=self.head_dim, n_heads=self.num_heads).unbind(dim=0)
+
+    def positional_embedding(self, q, k, v, kv_cache, layer_idx):
+        # Apply the rotary embeddings
+        seq_pos_offset = 0
+        if self.rotary_emb is not None and self.attention_axis == AttentionAxis.TIME:
+            if kv_cache is not None:
+                seq_pos_offset = kv_cache.seq_len(layer_idx)
+
+            # We need to permute because rotary embeddings expect the sequence dimension to be the second-to-last dimension
+            q, k = self.rotary_emb.rotate_queries_and_keys(q, k, seq_pos_offset=seq_pos_offset)
+
+        if kv_cache is not None and self.attention_axis == AttentionAxis.TIME:
+            # First, we append the current input key and value tensors to the cache.
+            # This concatenates the current key and value tensors to the existing key and value tensors
+            kv_cache.append(layer_idx, (k, v))
+            # Then, we retrieve the key and value tensors from the cache.
+            # This includes all the key and value tensors from previous time steps
+            # as well as the current time step.
+            k, v = kv_cache[layer_idx]
+
+        q = q.contiguous()
+        k = k.contiguous().to(q.dtype)  # Ensure k is the same dtype as q; this is necessary when using mixed precision
+        v = v.contiguous().to(q.dtype)  # Ensure v is the same dtype as q; this is necessary when using mixed precision
+
+        return q, k, v, seq_pos_offset
+
+    def rearrange_output(self, output: torch.Tensor, batch: int, variate: int, seq_len: int) -> torch.Tensor:
+        if self.attention_axis == AttentionAxis.TIME and self.use_memory_efficient_attention:
+            pattern = "(batch variate) seq_len n_heads head_dim -> batch variate seq_len (n_heads head_dim)"
+        elif self.attention_axis == AttentionAxis.TIME and not self.use_memory_efficient_attention:
+            pattern = "(batch variate) n_heads seq_len head_dim -> batch variate seq_len (n_heads head_dim)"
+        elif self.attention_axis == AttentionAxis.SPACE and self.use_memory_efficient_attention:
+            pattern = "(batch seq_len) variate n_heads head_dim -> batch variate seq_len (n_heads head_dim)"
+        elif self.attention_axis == AttentionAxis.SPACE and not self.use_memory_efficient_attention:
+            pattern = "(batch seq_len) n_heads variate head_dim -> batch variate seq_len (n_heads head_dim)"
+
+        return rearrange(output, pattern, batch=batch, variate=variate, seq_len=seq_len)  # type: ignore
+
+    def run_attention(self, attention_mask, q, k, v, seq_pos_offset, dropout, seq_len, variate):
+        # Determine dimension ranges for attention
+        # Ensure the last query vector index is used from the cache
+        q_dim_start, q_dim_end = seq_pos_offset, seq_pos_offset + seq_len
+        kv_dim_start, kv_dim_end = 0, v.shape[1] if self.use_memory_efficient_attention else v.shape[2]
+        if self.attention_axis == AttentionAxis.TIME:
+            attention_mask = (
+                attention_mask[..., q_dim_start:q_dim_end, kv_dim_start:kv_dim_end]
+                if torch.is_tensor(attention_mask)
+                else None
+            )
+            return scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=attention_mask,
+                dropout_p=dropout,
+                is_causal=(attention_mask is None and seq_pos_offset == 0),
+            )
+        elif self.attention_axis == AttentionAxis.SPACE:
+            # We don't use causal masking for space-wise attention
+            attention_mask = (
+                attention_mask[..., kv_dim_start:kv_dim_end, kv_dim_start:kv_dim_end]
+                if torch.is_tensor(attention_mask)
+                else None
+            )
+            return scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, dropout_p=dropout, is_causal=False)
+        else:
+            raise ValueError("Invalid attention axis")
+
+    def forward(
+        self,
+        layer_idx: int,
+        inputs: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        kv_cache=None,
+    ) -> torch.Tensor:
+        batch_size, variate, seq_len, _ = inputs.shape
+        dropout = self.dropout if self.training else 0.0
+
+        rearranged_inputs = self.rearrange_inputs(inputs)
+        q, k, v = self.get_qkv(rearranged_inputs)
+
+        q, k, v, seq_pos_offset = self.positional_embedding(q, k, v, kv_cache, layer_idx)
+
+        output = self.run_attention(attention_mask, q, k, v, seq_pos_offset, dropout, seq_len, variate)
+
+        output = self.rearrange_output(output, batch_size, variate, seq_len)
+        return self.wO(output)
+
+
+class TimeWiseMultiheadAttention(BaseMultiheadAttention):
+    """
+    Computes standard multihead causal attention over the time axis.
+    It does this by flattening out the variates along the batch dimension.
+    It also applies rotary position embeddings to the query and key matrices
+    in order to incorporate relative positional information.
+    """
+
+    attention_axis = AttentionAxis.TIME
+
+
+class SpaceWiseMultiheadAttention(BaseMultiheadAttention):
+    """
+    Computes bidirectional multihead attention over the space axis (i.e. across variates within
+    a multi-variate time series). This is done by flattening out the time axis along the batch dimension.
+    This allows the model to attend to different variates at the same time point. By alternating
+    between time-wise and space-wise attention, the model can learn both temporal and cross-variate
+    dependencies in the data.
+
+    Unlike with time-wise attention, don't apply rotary embeddings here
+    because we want cross-variate attention to be invariant to the order of the variates.
+    """
+
+    attention_axis = AttentionAxis.SPACE
+
+
+MultiHeadAttention = Union[TimeWiseMultiheadAttention, SpaceWiseMultiheadAttention]

autogluon.timeseries-1.4.1b20251011/src/autogluon/timeseries/models/toto/_internal/backbone/backbone.py

@@ -0,0 +1,262 @@
+# Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+#
+# This product includes software developed at Datadog (https://www.datadoghq.com/)
+# Copyright 2025 Datadog, Inc.
+
+import math
+from typing import NamedTuple, Optional
+
+import torch
+
+from .distribution import MixtureOfStudentTsOutput
+from .kvcache import KVCache
+from .scaler import CausalPatchStdMeanScaler
+from .transformer import Transformer
+
+
+class TotoOutput(NamedTuple):
+    """
+    Output of the Toto model. Contains the output distribution, the location parameters,
+    and the scale parameters.
+    """
+
+    distribution: torch.distributions.Distribution
+    loc: torch.Tensor
+    scale: torch.Tensor
+
+
+def patchify_id_mask(id_mask: torch.Tensor, patch_size: int) -> torch.Tensor:
+    patched_id_mask = id_mask.unfold(dimension=-1, size=patch_size, step=patch_size)
+    patched_id_mask_min = patched_id_mask.min(-1).values
+    patched_id_mask_max = patched_id_mask.max(-1).values
+    assert torch.eq(patched_id_mask_min, patched_id_mask_max).all(), "Patches cannot span multiple datasets"
+    return patched_id_mask_min
+
+
+class PatchEmbedding(torch.nn.Module):
+    """
+    Multivariate time series patch embedding.
+    Patchifies each variate separately.
+    """
+
+    def __init__(self, patch_size: int, stride: int, embed_dim: int):
+        super().__init__()
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.stride = stride
+        self.projection = torch.nn.Linear(self.patch_size, self.embed_dim)
+
+    def _patchify(self, x: torch.Tensor) -> torch.Tensor:
+        return x.unfold(dimension=-1, size=self.patch_size, step=self.stride)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        id_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert x.shape[-1] % self.patch_size == 0, (
+            f"Series length ({x.shape=}) must be divisible by ({self.patch_size=})"
+        )
+        x_patched: torch.Tensor = self._patchify(x)
+        id_mask_patched: torch.Tensor = self._patchify(id_mask)
+
+        assert torch.eq(id_mask_patched.min(-1).values, id_mask_patched.max(-1).values).all(), (
+            "Patches cannot span multiple datasets"
+        )
+
+        return (
+            self.projection(x_patched),
+            id_mask_patched.min(-1).values,
+        )
+
+
+class TotoBackbone(torch.nn.Module):
+    """
+    Toto (Timeseries-Optimized Transformer for Observability) is a transformer-based model for multivariate
+    time series forecasting. It applies a patch embedding to the input data, followed by a transformer
+    that alternates between time-wise and space-wise attention. The transformer is followed by a linear projection
+    that maps the transformer output to the output distribution.
+
+    The output distribution can be a single distribution (e.g. Gaussian) or a mixture of distributions.
+    If a mixture of distributions is used, the model will learn to predict the mixture weights
+    as well as the parameters of the individual distributions.
+
+    Parameters
+    ----------
+    patch_size
+        Size of the patch to use for the patch embedding.
+    stride
+        Stride to use for the patch embedding.
+    embed_dim
+        Dimension of the model's latent space.
+    num_layers
+        Number of transformer layers to use.
+    num_heads
+        Number of attention heads to use in each self-attention layer.
+    mlp_hidden_dim
+        Dimension of the hidden layer in the feedforward network.
+    dropout
+        Dropout rate to use in the model.
+    spacewise_every_n_layers
+        How many time-wise transformer layers to apply between each space-wise transformer layer.
+    spacewise_first
+        Whether to apply space-wise attention before time-wise attention.
+    scaler_cls
+        Class to use for scaling the input data.
+    output_distribution_classes
+        List of classes to use for the output distribution. If a single class is provided, the model
+        will output a single distribution. If multiple classes are provided, the model will output a
+        learned mixture of distributions.
+    output_distribution_kwargs
+        Keyword arguments to pass to the output distribution class. Note: this currently only works
+        with a single output distribution class.
+    use_memory_efficient_attention:
+        Whether to use memory-efficient attention. If True, the model will use the memory-efficient from xFormers.
+    stabilize_with_global:
+        Whether to use global statistics to stabilize causal statistics by clamping extreme values. Only applies to causal scalers.
+    scale_factor_exponent:
+        Exponent that controls the allowed range of deviation from global scale for causal scalers.
+    """
+
+    def __init__(
+        self,
+        patch_size: int,
+        stride: int,
+        embed_dim: int,
+        num_layers: int,
+        num_heads: int,
+        mlp_hidden_dim: int,
+        dropout: float,
+        spacewise_every_n_layers: int,
+        scaler_cls: str,
+        output_distribution_classes: list[str],
+        spacewise_first: bool = True,
+        output_distribution_kwargs: Optional[dict] = None,
+        use_memory_efficient_attention: bool = True,
+        stabilize_with_global: bool = True,
+        scale_factor_exponent: float = 10.0,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        # strings are used when loading a safetensors checkpoint
+        # Initialize patch-based scalers with the correct patch_size
+
+        self.scaler = CausalPatchStdMeanScaler(
+            patch_size=patch_size,
+            stabilize_with_global=stabilize_with_global,
+            scale_factor_exponent=scale_factor_exponent,
+        )
+        self.patch_embed = PatchEmbedding(patch_size, stride, embed_dim)
+        self.dropout = dropout
+        self.num_layers = num_layers
+        self.use_memory_efficient_attention = use_memory_efficient_attention
+        self.transformer = Transformer(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            num_layers=self.num_layers,
+            mlp_hidden_dim=mlp_hidden_dim,
+            dropout=dropout,
+            spacewise_every_n_layers=spacewise_every_n_layers,
+            spacewise_first=spacewise_first,
+            use_memory_efficient_attention=self.use_memory_efficient_attention,
+        )
+        self.unembed = torch.nn.Linear(embed_dim, embed_dim * patch_size)
+
+        # TODO[BEN] this doesn't need to be a list
+        output_distribution_classes_ = [MixtureOfStudentTsOutput]
+        self.output_distribution = output_distribution_classes_[0](embed_dim, **(output_distribution_kwargs or {}))
+
+    def allocate_kv_cache(
+        self,
+        batch_size: int,
+        num_variates: int,
+        max_time_steps: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> KVCache:
+        return KVCache(
+            batch_size=batch_size,
+            num_variates=num_variates,
+            transformer_layers=list(self.transformer.layers),
+            num_layers=self.num_layers,
+            embed_dim=self.embed_dim,
+            num_heads=self.transformer.layers[0].num_heads,  # type: ignore
+            max_seq_len=math.ceil(max_time_steps / self.patch_embed.stride),
+            device=device,
+            dtype=dtype,
+            use_memory_efficient_attention=self.use_memory_efficient_attention,
+        )
+
+    def backbone(
+        self,
+        inputs: torch.Tensor,
+        input_padding_mask: torch.Tensor,
+        id_mask: torch.Tensor,
+        kv_cache: Optional[KVCache] = None,
+        scaling_prefix_length: Optional[int] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        scaled_inputs: torch.Tensor
+        loc: torch.Tensor
+        scale: torch.Tensor
+
+        # Standard scaling operation, same API but without ID mask.
+        scaled_inputs, loc, scale = self.scaler(
+            inputs,
+            weights=torch.ones_like(inputs, device=inputs.device),
+            padding_mask=input_padding_mask,
+            prefix_length=scaling_prefix_length,
+        )
+
+        if kv_cache is not None:
+            prefix_len = self.patch_embed.stride * kv_cache.current_len(0)
+
+            # Truncate inputs so that the transformer only processes
+            # the last patch in the sequence. We'll use the KVCache
+            # for the earlier patches.
+            scaled_inputs = scaled_inputs[:, :, prefix_len:]
+
+            # As a simplification, when using kv cache we only allow decoding
+            # one step at a time after the initial forward pass.
+            assert (prefix_len == 0) or (scaled_inputs.shape[-1] == self.patch_embed.stride), (
+                "Must decode one step at a time."
+            )
+
+            input_padding_mask = input_padding_mask[:, :, prefix_len:]
+            id_mask = id_mask[:, :, prefix_len:]
+
+        embeddings: torch.Tensor
+        reduced_id_mask: torch.Tensor
+
+        embeddings, reduced_id_mask = self.patch_embed(scaled_inputs, id_mask)
+
+        # Apply the transformer on the embeddings
+        transformed: torch.Tensor = self.transformer(embeddings, reduced_id_mask, kv_cache)
+
+        # Unembed and flatten the sequence
+        unembedded = self.unembed(transformed)
+        batch_size, num_variates, seq_len = unembedded.shape[:3]
+        patch_size = unembedded.shape[-1] // self.embed_dim
+        flattened = unembedded.view(batch_size, num_variates, seq_len * patch_size, self.embed_dim)
+        return flattened, loc, scale
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        input_padding_mask: torch.Tensor,
+        id_mask: torch.Tensor,
+        kv_cache: Optional[KVCache] = None,
+        scaling_prefix_length: Optional[int] = None,
+    ) -> TotoOutput:
+        flattened, loc, scale = self.backbone(
+            inputs,
+            input_padding_mask,
+            id_mask,
+            kv_cache,
+            scaling_prefix_length,
+        )
+
+        return TotoOutput(self.output_distribution(flattened), loc, scale)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device

autogluon.timeseries-1.4.1b20251011/src/autogluon/timeseries/models/toto/_internal/backbone/distribution.py

@@ -0,0 +1,70 @@
+# Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License.
+#
+# This product includes software developed at Datadog (https://www.datadoghq.com/)
+# Copyright 2025 Datadog, Inc.
+
+from abc import ABC
+
+import torch
+import torch.nn.functional as F
+from gluonts.torch.distributions import AffineTransformed
+from gluonts.torch.distributions.studentT import StudentT
+
+
+class DistributionOutput(ABC, torch.nn.Module):
+    pass
+
+
+class StudentTOutput(DistributionOutput):
+    def __init__(self, embed_dim):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.df = torch.nn.Linear(embed_dim, 1)
+        self.loc_proj = torch.nn.Linear(embed_dim, 1)
+        self.scale_proj = torch.nn.Linear(embed_dim, 1)
+
+    def forward(self, inputs, loc=None, scale=None):
+        eps = torch.finfo(inputs.dtype).eps
+        df = 2.0 + F.softplus(self.df(inputs)).clamp_min(eps).squeeze(-1)
+        base_loc = self.loc_proj(inputs).squeeze(-1)
+        base_scale = F.softplus(self.scale_proj(inputs)).clamp_min(eps).squeeze(-1)
+
+        base_dist = torch.distributions.StudentT(df, base_loc, base_scale, validate_args=False)  # type: ignore
+
+        if loc is not None and scale is not None:
+            return AffineTransformed(
+                base_dist,
+                loc=loc,
+                scale=scale,
+            )
+        return base_dist
+
+
+class MixtureOfStudentTsOutput(DistributionOutput):
+    def __init__(
+        self,
+        embed_dim,
+        k_components,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.k_components = k_components
+
+        self.df = torch.nn.Linear(embed_dim, k_components)
+        self.loc_proj = torch.nn.Linear(embed_dim, k_components)
+        self.scale_proj = torch.nn.Linear(embed_dim, k_components)
+        self.mixture_weights = torch.nn.Linear(embed_dim, k_components)
+
+    def forward(self, inputs, loc=None, scale=None):
+        df = 2.0 + F.softplus(self.df(inputs)).clamp_min(torch.finfo(inputs.dtype).eps)
+        loc = self.loc_proj(inputs)
+        scale = F.softplus(self.scale_proj(inputs)).clamp_min(torch.finfo(inputs.dtype).eps)
+        logits = self.mixture_weights(inputs)
+        probs = F.softmax(logits, dim=-1)
+        components = StudentT(df, loc, scale)
+        mixture_distribution = torch.distributions.Categorical(probs=probs)
+
+        return torch.distributions.MixtureSameFamily(
+            mixture_distribution,
+            components,
+        )