dtw_loss_functions.soft_dtw_implementations.soft_dtw_cuda_ron.autograd — dtw_loss_functions documentation

# Source code for dtw_loss_functions.soft_dtw_implementations.soft_dtw_cuda_ron.autograd

    from __future__ import annotations

    import torch
    from torch.autograd import Function

    from .utils.checks import check_D
    from .cuda.launcher import softdtw_forward_cuda, softdtw_backward_cuda_log
    from .cuda.launcher import softdtw_forward_cpu, softdtw_backward_cpu

    class SoftDTWAutograd(Function):

        @staticmethod
        def forward(ctx, D: torch.Tensor, gamma: float, bandwidth: float | None):
            check_D(D)
            gamma_f = float(gamma)
            if gamma_f <= 0:
                raise ValueError(f"gamma must be > 0, got {gamma_f}")

            if bandwidth is None:
                bandwidth_f = -1.0
            else:
                bw = float(bandwidth)
                bandwidth_f = -1.0 if bw <= 0 else bw

            if D.is_cuda:
                out, R = softdtw_forward_cuda(D, gamma_f, bandwidth_f)
            else:
                out, R = softdtw_forward_cpu(D, gamma_f, bandwidth_f)

            ctx.save_for_backward(D, R.detach())
            ctx.gamma = gamma_f
            ctx.bandwidth = bandwidth_f
            return out

        @staticmethod
        def backward(ctx, grad_output: torch.Tensor):
            D, R = ctx.saved_tensors
            gamma_f = ctx.gamma
            bandwidth_f = ctx.bandwidth

            if D.is_cuda:
                E = softdtw_backward_cuda_log(D, R, gamma_f, bandwidth_f)
            else:
                E = softdtw_backward_cpu(D, R, gamma_f, bandwidth_f)

            g = grad_output.reshape(-1).to(dtype=E.dtype).view(-1, 1, 1)
            grad_D = g * E
            return grad_D, None, None