transport.samples.kernels.grad_sqeuclid.sinkhorn_online_grad_sqeuclid
def sinkhorn_online_grad_sqeuclid(x: torch.Tensor, y: torch.Tensor, a: torch.Tensor, b: torch.Tensor, f: torch.Tensor, g: torch.Tensor, eps: float, allow_tf32: bool = True, use_exp2: bool = True, block_m: int | None = None, block_n: int | None = None, block_k: int | None = None, num_warps: int | None = None, num_stages: int = 2, autotune: bool = True, grad_scale: torch.Tensor | None = None, compute_grad_x: bool = True, compute_grad_y: bool = True, cost_scale: float = 1.0, label_x: torch.Tensor | None = None, label_y: torch.Tensor | None = None, label_cost_matrix: torch.Tensor | None = None, lambda_x: float = 1.0, lambda_y: float = 0.0) -> tuple[torch.Tensor, torch.Tensor]Source: torchmatch/transport/samples/kernels/grad_sqeuclid.py:941