ml4gw 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

ml4gw/transforms/spectral.py

@@ -1,8 +1,11 @@
 from typing import Optional
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 from ml4gw.spectral import fast_spectral_density, spectral_density
+from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d
 
 
 class SpectralDensity(torch.nn.Module):
@@ -51,7 +54,9 @@ class SpectralDensity(torch.nn.Module):
         fftlength: float,
         overlap: Optional[float] = None,
         average: str = "mean",
-        window: Optional[torch.Tensor] = None,
+        window: Optional[
+            Float[Tensor, " {int(fftlength*sample_rate)}"]
+        ] = None,
         fast: bool = False,
     ) -> None:
         if overlap is None:
@@ -93,7 +98,9 @@ class SpectralDensity(torch.nn.Module):
         self.average = average
         self.fast = fast
 
-    def forward(self, x: torch.Tensor, y: Optional[torch.Tensor] = None):
+    def forward(
+        self, x: TimeSeries1to3d, y: Optional[TimeSeries1to3d] = None
+    ) -> FrequencySeries1to3d:
         if self.fast:
             return fast_spectral_density(
                 x,

ml4gw/transforms/spectrogram.py

@@ -3,8 +3,12 @@ from typing import Dict, List
 
 import torch
 import torch.nn.functional as F
+from jaxtyping import Float
+from torch import Tensor
 from torchaudio.transforms import Spectrogram
 
+from ml4gw.types import TimeSeries3d
+
 
 class MultiResolutionSpectrogram(torch.nn.Module):
     """
@@ -122,7 +126,9 @@ class MultiResolutionSpectrogram(torch.nn.Module):
 
         return [dict(zip(kwargs, col)) for col in zip(*kwargs.values())]
 
-    def forward(self, X: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, X: TimeSeries3d
+    ) -> Float[Tensor, "batch channel frequency time"]:
         """
         Calculate spectrograms of the input tensor and
         combine them into a single spectrogram

ml4gw/transforms/spline_interpolation.py

@@ -0,0 +1,370 @@
+"""
+Adaptation of code from https://github.com/dottormale/Qtransform
+"""
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class SplineInterpolate(torch.nn.Module):
+    """
+    Perform 1D or 2D spline interpolation based on De Boor's method.
+    Supports batched, multi-channel inputs, so acceptable data
+    shapes are `(width)`, `(height, width)`, `(batch, width)`,
+    `(batch, height, width)`, `(batch, channel, width)`, and
+    `(batch, channel, height, width)`.
+
+    During initialization of this Module, both the desired input
+    and output coordinate Tensors can be specified to allow
+    pre-computation of the B-spline basis matrices, though the only
+    mandatory argument is the coordinates of the data along the
+    `width` dimension. If no argument is given for coordinates along
+    the `height` dimension, it is assumed that 1D interpolation is
+    desired.
+
+    Unlike scipy's implementation of spline interpolation, the data
+    to be interpolated is not passed until actually calling the
+    object. This is useful for cases where the input and output
+    coordinates are known in advance, but the data is not, so that
+    the interpolator can be set up ahead of time.
+
+    WARNING: compared to scipy's spline interpolation, this method
+    produces edge artifacts when the output coordinates are near
+    the boundaries of the input coordinates. Therefore, it is
+    recommended to interpolate only to coordinates that are well
+    within the input coordinate range. Unfortunately, the specific
+    definition of "well within" changes based on the size of the
+    data, so some testing may be required to get good results.
+
+    Args:
+        x_in:
+            Coordinates of the width dimension of the data
+        y_in:
+            Coordinates of the height dimension of the data. If not
+            specified, it is assumed the 1D interpolation is desired,
+            and so the default value is a Tensor of length 1
+        kx:
+            Degree of spline interpolation along the width dimension.
+            Default is cubic.
+        ky:
+            Degree of spline interpolation along the height dimension.
+            Default is cubic.
+        sx:
+            Regularization factor to avoid singularities during matrix
+            inversion for interpolation along the width dimension. Not
+            to be confused with the `s` parameter in scipy's spline
+            methods, which controls the number of knots.
+        sy:
+            Regularization factor to avoid singularities during matrix
+            inversion for interpolation along the height dimension.
+        x_out:
+            Coordinates for the data to be interpolated to along the
+            width dimension. If not specified during initialization,
+            this must be specified during the object call.
+        y_out:
+            Coordinates for the data to be interpolated to along the
+            height dimension. If not specified during initialization,
+            this must be specified during the object call.
+
+    """
+
+    def __init__(
+        self,
+        x_in: Tensor,
+        y_in: Tensor = Tensor([1]),
+        kx: int = 3,
+        ky: int = 3,
+        sx: float = 0.001,
+        sy: float = 0.001,
+        x_out: Optional[Tensor] = None,
+        y_out: Optional[Tensor] = None,
+    ):
+        super().__init__()
+        self.kx = kx
+        self.ky = ky
+        self.sx = sx
+        self.sy = sy
+        self.register_buffer("x_in", x_in)
+        self.register_buffer("y_in", y_in)
+        self.register_buffer("x_out", x_out)
+        self.register_buffer("y_out", y_out)
+
+        tx, Bx, BxT_Bx = self._compute_knots_and_basis_matrices(x_in, kx, sx)
+        self.register_buffer("tx", tx)
+        self.register_buffer("Bx", Bx)
+        self.register_buffer("BxT_Bx", BxT_Bx)
+
+        ty, By, ByT_By = self._compute_knots_and_basis_matrices(y_in, ky, sy)
+        self.register_buffer("ty", ty)
+        self.register_buffer("By", By)
+        self.register_buffer("ByT_By", ByT_By)
+
+        if self.x_out is not None:
+            Bx_out = self.bspline_basis_natural(x_out, kx, self.tx)
+            self.register_buffer("Bx_out", Bx_out)
+        if self.y_out is not None:
+            By_out = self.bspline_basis_natural(y_out, ky, self.ty)
+            self.register_buffer("By_out", By_out)
+
+    def _compute_knots_and_basis_matrices(self, x, k, s):
+        knots = self.generate_natural_knots(x, k)
+        basis_matrix = self.bspline_basis_natural(x, k, knots)
+        identity = torch.eye(basis_matrix.shape[-1])
+        B_T_B = basis_matrix.T @ basis_matrix + s * identity
+        return knots, basis_matrix, B_T_B
+
+    def generate_natural_knots(self, x: Tensor, k: int) -> Tensor:
+        """
+        Generates a natural knot sequence for B-spline interpolation.
+        Natural knot sequence means that 2*k knots are added to the beginning
+        and end of datapoints as replicas of first and last datapoint
+        respectively in order to enforce natural boundary conditions,
+        i.e. second derivative = 0.
+        The other n nodes are placed in correspondece of the data points.
+
+        Args:
+            x: Tensor of data point positions.
+            k: Degree of the spline.
+
+        Returns:
+            Tensor of knot positions.
+        """
+        return F.pad(x[None], (k, k), mode="replicate")[0]
+
+    def compute_L_R(
+        self,
+        x: Tensor,
+        t: Tensor,
+        d: int,
+        m: int,
+    ) -> Tuple[Tensor, Tensor]:
+
+        """
+        Compute the L and R values for B-spline basis functions.
+        L and R are respectively the first and second coefficient multiplying
+        B_{i,p-1}(x) and B_{i+1,p-1}(x) in De Boor's recursive formula for
+        Bspline basis funciton computation
+        See https://en.wikipedia.org/wiki/De_Boor%27s_algorithm for details
+
+        Args:
+            x:
+                Tensor of data point positions.
+            t:
+                Tensor of knot positions.
+            d:
+                Current degree of the basis function.
+            m:
+                Number of intervals (n - k - 1, where n is the number of knots
+                and k is the degree).
+
+        Returns:
+            L: Tensor containing left values for the B-spline basis functions.
+            R: Tensor containing right values for the B-spline basis functions.
+        """
+        left_num = x.unsqueeze(1) - t[:m].unsqueeze(0)
+        left_den = t[d : m + d] - t[:m]
+        L = left_num / left_den.unsqueeze(0)
+        L = torch.nan_to_num_(L, nan=0.0, posinf=0.0, neginf=0.0)
+
+        right_num = t[d + 1 : m + d + 1] - x.unsqueeze(1)
+        right_den = t[d + 1 : m + d + 1] - t[1 : m + 1]
+        R = right_num / right_den.unsqueeze(0)
+        R = torch.nan_to_num_(R, nan=0.0, posinf=0.0, neginf=0.0)
+
+        return L, R
+
+    def zeroth_order(
+        self,
+        x: Tensor,
+        k: int,
+        t: Tensor,
+        n: int,
+        m: int,
+    ) -> Tensor:
+
+        """
+        Compute the zeroth-order B-spline basis functions
+        according to de Boors recursive formula.
+        See https://en.wikipedia.org/wiki/De_Boor%27s_algorithm for reference
+
+        Args:
+            x:
+                Tensor of data point positions.
+            k:
+                Degree of the spline.
+            t:
+                Tensor of knot positions.
+            n:
+                Number of data points.
+            m:
+                Number of intervals (n - k - 1, where n is the number of knots
+                and k is the degree).
+
+        Returns:
+            b: Tensor containing the zeroth-order B-spline basis functions.
+        """
+        b = torch.zeros((n, m, k + 1))
+
+        mask_lower = t[: m + 1].unsqueeze(0)[:, :-1] <= x.unsqueeze(1)
+        mask_upper = x.unsqueeze(1) < t[: m + 1].unsqueeze(0)[:, 1:]
+
+        b[:, :, 0] = mask_lower & mask_upper
+        b[:, 0, 0] = torch.where(x < t[1], torch.ones_like(x), b[:, 0, 0])
+        b[:, -1, 0] = torch.where(x >= t[-2], torch.ones_like(x), b[:, -1, 0])
+        return b
+
+    def bspline_basis_natural(
+        self,
+        x: Tensor,
+        k: int,
+        t: Tensor,
+    ) -> Tensor:
+        """
+        Compute bspline basis function using de Boor's recursive formula
+        (See https://en.wikipedia.org/wiki/De_Boor%27s_algorithm for reference)
+        Args:
+            x: Tensor of data point positions.
+            k: Degree of the spline.
+            t: Tensor of knot positions.
+
+        Returns:
+            Tensor containing the kth-order B-spline basis functions
+        """
+
+        if len(x) == 1:
+            return torch.eye(1)
+        n = x.shape[0]
+        m = t.shape[0] - k - 1
+
+        # calculate zeroth order basis funciton
+        b = self.zeroth_order(x, k, t, n, m)
+
+        zeros_tensor = torch.zeros(b.shape[0], 1)
+        # recursive de Boors formula for bspline basis functions
+        for d in range(1, k + 1):
+            L, R = self.compute_L_R(x, t, d, m)
+            left = L * b[:, :, d - 1]
+
+            temp_b = torch.cat([b[:, 1:, d - 1], zeros_tensor], dim=1)
+
+            right = R * temp_b
+            b[:, :, d] = left + right
+
+        return b[:, :, -1]
+
+    def bivariate_spline_fit_natural(self, Z):
+
+        if len(Z.shape) == 3:
+            Z_Bx = torch.matmul(Z, self.Bx)
+            # ((BxT @ Bx)^-1 @ (Z @ Bx)T)T = Z @ BxT^-1
+            return torch.linalg.solve(self.BxT_Bx, Z_Bx.mT).mT
+
+        # Adding batch/channel dimension handling
+        # ByT @ Z @ Bx
+        ByT_Z_Bx = torch.einsum("ij,bcik,kl->bcjl", self.By, Z, self.Bx)
+        # (ByT @ By)^-1 @ (ByT @ Z @ Bx) = By^-1 @ Z @ Bx
+        E = torch.linalg.solve(self.ByT_By, ByT_Z_Bx)
+        # ((BxT @ Bx)^-1 @ (By^-1 @ Z @ Bx)T)T = By^-1 @ Z @ BxT^-1
+        return torch.linalg.solve(self.BxT_Bx, E.mT).mT
+
+    def evaluate_bivariate_spline(self, C: Tensor):
+        """
+        Evaluate a bivariate spline on a grid of x and y points.
+
+        Args:
+            C: Coefficient tensor of shape (batch_size, mx, my).
+
+        Returns:
+            Z_interp: Interpolated values at the grid points.
+        """
+        # Perform matrix multiplication using einsum to get Z_interp
+        if len(C.shape) == 3:
+            return torch.matmul(C, self.Bx_out.mT)
+        return torch.einsum("ik,bckm,mj->bcij", self.By_out, C, self.Bx_out.mT)
+
+    def _validate_inputs(self, Z, x_out, y_out):
+        if x_out is None and self.x_out is None:
+            raise ValueError(
+                "Output x-coordinates were not specified in either object "
+                "creation or in forward call"
+            )
+
+        if y_out is None and self.y_out is None:
+            y_out = self.y_in
+
+        dims = len(Z.shape)
+        if dims > 4:
+            raise ValueError("Input data has more than 4 dimensions")
+
+        if len(self.y_in) > 1 and dims == 1:
+            raise ValueError(
+                "An input y-coordinate array with length greater than 1 "
+                "was given, but the input data is 1-dimensional. Expected "
+                "input data to be at least 2-dimensional"
+            )
+
+        # Expand Z to have 4 dimensions
+        # There are 6 valid input shapes: (w), (b, w), (b, c, w),
+        # (h, w), (b, h, w), and (b, c, h, w).
+
+        # If the input y coordinate array has length 1,
+        # assume the first dimension(s) are batch dimensions
+        # and that no height dimension is included in Z
+        idx = -2 if len(self.y_in) == 1 else -3
+        while len(Z.shape) < 4:
+            Z = Z.unsqueeze(idx)
+
+        if Z.shape[-2:] != torch.Size([len(self.y_in), len(self.x_in)]):
+            raise ValueError(
+                "The spatial dimensions of the data tensor do not match "
+                "the given input dimensions. "
+                f"Expected [{len(self.y_in)}, {len(self.x_in)}], but got "
+                f"[{Z.shape[-2]}, {Z.shape[-1]}]"
+            )
+
+        return Z, y_out
+
+    def forward(
+        self,
+        Z: Tensor,
+        x_out: Optional[Tensor] = None,
+        y_out: Optional[Tensor] = None,
+    ) -> Tensor:
+        """
+        Compute the interpolated data
+
+        Args:
+            Z:
+                Tensor of data to be interpolated. Must be between 1 and 4
+                dimensions. The shape of the tensor must agree with the
+                input coordinates given on initialization. If `y_in` was
+                not specified during initialization, it is assumed that
+                Z does not have a height dimension.
+            x_out:
+                Coordinates to interpolate the data to along the width
+                dimension. Overrides any value that was set during
+                initialization.
+            y_out:
+                Coordinates to interpolate the data to along the height
+                dimension. Overrides any value that was set during
+                initialization.
+
+        Returns:
+            A 4D tensor with shape `(batch, channel, height, width)`.
+            Depending on the input data shape, many of these dimensions
+            may have length 1.
+        """
+
+        Z, y_out = self._validate_inputs(Z, x_out, y_out)
+
+        if x_out is not None:
+            self.Bx_out = self.bspline_basis_natural(x_out, self.kx, self.tx)
+        if y_out is not None:
+            self.By_out = self.bspline_basis_natural(y_out, self.ky, self.ty)
+
+        coef = self.bivariate_spline_fit_natural(Z)
+        Z_interp = self.evaluate_bivariate_spline(coef)
+        return Z_interp

ml4gw/transforms/transform.py

@@ -3,6 +3,7 @@ from typing import Optional
 import torch
 
 from ml4gw.spectral import spectral_density
+from ml4gw.types import FrequencySeries1to3d, TimeSeries1to3d
 
 
 class FittableTransform(torch.nn.Module):
@@ -43,12 +44,12 @@ class FittableTransform(torch.nn.Module):
 class FittableSpectralTransform(FittableTransform):
     def normalize_psd(
         self,
-        x,
+        x: TimeSeries1to3d,
         sample_rate: float,
         num_freqs: int,
         fftlength: Optional[float] = None,
         overlap: Optional[float] = None,
-    ):
+    ) -> FrequencySeries1to3d:
         # if we specified an FFT length, convert
         # the (assumed) time-domain data to the
         # frequency domain
@@ -68,7 +69,7 @@ class FittableSpectralTransform(FittableTransform):
                 scale=scale,
             )
 
-        # add two dummy dimensions in case we need to inerpolate
+        # add two dummy dimensions in case we need to interpolate
         # the frequency dimension, since `interpolate` expects
         # a (batch, channel, spatial) formatted tensor as input
         x = x.view(1, 1, -1)

ml4gw/transforms/waveforms.py

@@ -1,8 +1,11 @@
 from typing import List, Optional
 
 import torch
+from jaxtyping import Float
+from torch import Tensor
 
 from ml4gw import gw
+from ml4gw.types import BatchTensor
 
 
 # TODO: should these live in ml4gw.waveforms submodule?
@@ -10,8 +13,8 @@ from ml4gw import gw
 class WaveformSampler(torch.nn.Module):
     def __init__(
         self,
-        parameters: Optional[torch.Tensor] = None,
-        **polarizations: torch.Tensor,
+        parameters: Optional[Float[Tensor, "batch num_params"]] = None,
+        **polarizations: Float[Tensor, "batch time"],
     ):
         super().__init__()
         # make sure we have the same number of waveforms
@@ -29,7 +32,7 @@ class WaveformSampler(torch.nn.Module):
             elif num_waveforms is None:
                 num_waveforms = tensor.shape[0]
 
-            self.polarizations[polarization] = torch.Tensor(tensor)
+            self.polarizations[polarization] = Tensor(tensor)
 
         if parameters is not None and len(parameters) != num_waveforms:
             raise ValueError(
@@ -73,10 +76,10 @@ class WaveformProjector(torch.nn.Module):
 
     def forward(
         self,
-        dec: gw.ScalarTensor,
-        psi: gw.ScalarTensor,
-        phi: gw.ScalarTensor,
-        **polarizations,
+        dec: BatchTensor,
+        psi: BatchTensor,
+        phi: BatchTensor,
+        **polarizations: Float[Tensor, "batch time"],
     ):
         ifo_responses = gw.compute_observed_strain(
             dec,

ml4gw/transforms/whitening.py

@@ -1,9 +1,15 @@
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 
 from ml4gw import spectral
 from ml4gw.transforms.transform import FittableSpectralTransform
+from ml4gw.types import (
+    FrequencySeries1d,
+    FrequencySeries1to3d,
+    TimeSeries1d,
+    TimeSeries3d,
+)
 
 
 class Whiten(torch.nn.Module):
@@ -58,7 +64,9 @@ class Whiten(torch.nn.Module):
         window = torch.hann_window(size, dtype=torch.float64)
         self.register_buffer("window", window)
 
-    def forward(self, X: torch.Tensor, psd: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, X: TimeSeries3d, psd: FrequencySeries1to3d
+    ) -> TimeSeries3d:
         """
         Whiten a batch of multichannel timeseries by a
         background power spectral density.
@@ -142,7 +150,7 @@ class FixedWhiten(FittableSpectralTransform):
     def fit(
         self,
         fduration: float,
-        *background: torch.Tensor,
+        *background: Union[TimeSeries1d, FrequencySeries1d],
         fftlength: Optional[float] = None,
         highpass: Optional[float] = None,
         overlap: Optional[float] = None
@@ -224,7 +232,7 @@ class FixedWhiten(FittableSpectralTransform):
         fduration = torch.Tensor([fduration])
         self.build(psd=psd, fduration=fduration)
 
-    def forward(self, X: torch.Tensor) -> torch.Tensor:
+    def forward(self, X: TimeSeries3d) -> TimeSeries3d:
         """
         Whiten the input timeseries tensor using the
         PSD fit by the `.fit` method, which must be

ml4gw/types.py

@@ -1,10 +1,25 @@
-from torchtyping import TensorType
-
-WaveformTensor = TensorType["batch", "num_ifos", "time"]
-PSDTensor = TensorType["num_ifos", "frequency"]
-ScalarTensor = TensorType["batch"]
-VectorGeometry = TensorType["batch", "space"]
-TensorGeometry = TensorType["batch", "space", "space"]
-NetworkVertices = TensorType["num_ifos", 3]
-NetworkDetectorTensors = TensorType["num_ifos", 3, 3]
-TimeSeriesTensor = TensorType["num_channels", "time"]
+from typing import Union
+
+from jaxtyping import Float
+from torch import Tensor
+
+WaveformTensor = Float[Tensor, "batch num_ifos time"]
+PSDTensor = Float[Tensor, "num_ifos frequency"]
+BatchTensor = Float[Tensor, "batch"]
+VectorGeometry = Float[Tensor, "batch space"]
+TensorGeometry = Float[Tensor, "batch space space"]
+NetworkVertices = Float[Tensor, "num_ifos 3"]
+NetworkDetectorTensors = Float[Tensor, "num_ifos 3 3"]
+
+
+TimeSeries1d = Float[Tensor, "time"]
+TimeSeries2d = Float[TimeSeries1d, "channel"]
+TimeSeries3d = Float[TimeSeries2d, "batch"]
+TimeSeries1to3d = Union[TimeSeries1d, TimeSeries2d, TimeSeries3d]
+
+FrequencySeries1d = Float[Tensor, "frequency"]
+FrequencySeries2d = Float[FrequencySeries1d, "channel"]
+FrequencySeries3d = Float[FrequencySeries2d, "batch"]
+FrequencySeries1to3d = Union[
+    FrequencySeries1d, FrequencySeries2d, FrequencySeries3d
+]

ml4gw/utils/interferometer.py

@@ -4,7 +4,7 @@ import torch
 # based on values from
 # https://lscsoft.docs.ligo.org/lalsuite/lal/_l_a_l_detectors_8h_source.html
 class InterferometerGeometry:
-    def __init__(self, name: str):
+    def __init__(self, name: str) -> None:
         if name == "H1":
             self.x_arm = torch.Tensor(
                 (-0.22389266154, +0.79983062746, +0.55690487831)

ml4gw/utils/slicing.py

@@ -1,25 +1,30 @@
 from typing import Optional, Union
 
 import torch
+from jaxtyping import Float, Int64
+from torch import Tensor
 from torch.nn.functional import unfold
-from torchtyping import TensorType
 
-# need to define these for flake8 compatibility
-batch = time = channel = None  # noqa
+from ml4gw.types import (
+    TimeSeries1d,
+    TimeSeries1to3d,
+    TimeSeries2d,
+    TimeSeries3d,
+)
 
-TimeSeriesTensor = Union[TensorType["time"], TensorType["channel", "time"]]
-
-BatchTimeSeriesTensor = Union[
-    TensorType["batch", "time"], TensorType["batch", "channel", "time"]
-]
+BatchTimeSeriesTensor = Union[Float[Tensor, "batch time"], TimeSeries3d]
 
 
 def unfold_windows(
-    x: torch.Tensor,
+    x: TimeSeries1to3d,
     window_size: int,
     stride: int,
     drop_last: bool = True,
-):
+) -> Union[
+    Float[TimeSeries1d, " window"],
+    Float[TimeSeries2d, " window"],
+    Float[TimeSeries3d, " window"],
+]:
     """Unfold a timeseries into windows
 
     Args:
@@ -83,8 +88,8 @@ def unfold_windows(
 
 
 def slice_kernels(
-    x: Union[TimeSeriesTensor, TensorType["batch", "channel", "time"]],
-    idx: TensorType[..., torch.int64],
+    x: TimeSeries1to3d,
+    idx: Int64[Tensor, "..."],
     kernel_size: int,
 ) -> BatchTimeSeriesTensor:
     """Slice kernels from single or multichannel timeseries
@@ -96,7 +101,8 @@ def slice_kernels(
     one more dimension than `x`.
 
     Args:
-        x: The timeseries tensor to slice kernels from
+        x:
+            The timeseries tensor to slice kernels from
         idx:
             The indices in `x` of the first sample of each
             kernel. If `x` is 1D, `idx` must be 1D as well.
@@ -114,6 +120,7 @@ def slice_kernels(
             coincidentally among the channels.
         kernel_size:
             The length of the kernels to slice from the timeseries
+
     Returns:
         A tensor of shape `(batch_size, kernel_size)` if `x` is
         1D and `(batch_size, num_channels, kernel_size)` if `x`
@@ -225,7 +232,7 @@ def slice_kernels(
 
 
 def sample_kernels(
-    X: TimeSeriesTensor,
+    X: TimeSeries1to3d,
     kernel_size: int,
     N: Optional[int] = None,
     max_center_offset: Optional[int] = None,
@@ -245,8 +252,9 @@ def sample_kernels(
     either be `None` or be equal to `len(X)`.
 
     Args:
-        X: The timeseries tensor from which to sample kernels
-        kernel_size: The size of the kernels to sample
+        X:
+            The timeseries tensor from which to sample kernels
+            kernel_size: The size of the kernels to sample
         N:
             The number of kernels to sample. Can be left as
             `None` if `X` is 3D, otherwise must be specified