dtw_loss_functions.soft_dtw_implementations.soft_dtw_cuda_ron.distances — dtw_loss_functions documentation

# Source code for dtw_loss_functions.soft_dtw_implementations.soft_dtw_cuda_ron.distances

    from __future__ import annotations

    import torch

    def sqeuclidean(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Efficient squared Euclidean distance:
          D[b,i,j] = ||x[b,i]-y[b,j]||^2

        x: (B,N,D), y: (B,M,D)
        returns: (B,N,M)
        """
        if x.dim() != 3 or y.dim() != 3:
            raise ValueError(f"Expected x,y as (B,N,D)/(B,M,D). Got {tuple(x.shape)} and {tuple(y.shape)}")
        if x.shape[0] != y.shape[0] or x.shape[2] != y.shape[2]:
            raise ValueError(f"Batch/features mismatch: {tuple(x.shape)} vs {tuple(y.shape)}")

        # (B,N)
        x2 = (x * x).sum(dim=-1)
        # (B,M)
        y2 = (y * y).sum(dim=-1)
        # (B,N,M)
        xy = torch.bmm(x, y.transpose(1, 2))
        D = x2.unsqueeze(2) + y2.unsqueeze(1) - 2.0 * xy

        # Numerical cleanup (fp roundoff can produce tiny negatives)
        return D.clamp_min(0.0)

    def pairwise_distance(x: torch.Tensor, y: torch.Tensor, *, dist: str) -> torch.Tensor:
        dist = dist.lower()
        if dist in ("sqeuclidean", "sq_euclidean", "squared_euclidean"):
            return sqeuclidean(x, y)
        raise ValueError(f"Unknown dist='{dist}'. Supported: 'sqeuclidean'")