Awesome!

Just for completeness, here is a self contained version of our current distance function in gpytorch, which is about 2x faster than the original snippet I posted in that issue:

def fast_cdist(x1, x2):
    adjustment = x1.mean(-2, keepdim=True)
    x1 = x1 - adjustment
    x2 = x2 - adjustment  # x1 and x2 should be identical in all dims except -2 at this point

    # Compute squared distance matrix using quadratic expansion
    # But be clever and do it with a single matmul call
    x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
    x1_pad = torch.ones_like(x1_norm)
    x2_norm = x2.pow(2).sum(dim=-1, keepdim=True)
    x2_pad = torch.ones_like(x2_norm)
    x1_ = torch.cat([-2. * x1, x1_norm, x1_pad], dim=-1)
    x2_ = torch.cat([x2, x2_pad, x2_norm], dim=-1)
    res = x1_.matmul(x2_.transpose(-2, -1))

    # Zero out negative values
    res.clamp_min_(1e-30).sqrt_()
    return res

This uses the same basic strategy as the original method I'd written, but does everything in a single matmul call because matmul calls are really just that fast.

That being said, it has some drawbacks for inclusion upstream in pytorch proper over my original snippet. Namely (1) the additional working memory overhead, and (2) it's the same basic strategy, so a proper implementation of the first snippet I'd linked may erase the gap anyways.

matmul

Awesome!

Just for completeness, here is a self contained version of our current distance function in gpytorch, which is about 2x faster than the original snippet I posted in that issue:

def fast_cdist(x1, x2):
    adjustment = x1.mean(-2, keepdim=True)
    x1 = x1 - adjustment
    x2 = x2 - adjustment  # x1 and x2 should be identical in all dims except -2 at this point

    # Compute squared distance matrix using quadratic expansion
    # But be clever and do it with a single matmul call
    x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
    x1_pad = torch.ones_like(x1_norm)
    x2_norm = x2.pow(2).sum(dim=-1, keepdim=True)
    x2_pad = torch.ones_like(x2_norm)
    x1_ = torch.cat([-2. * x1, x1_norm, x1_pad], dim=-1)
    x2_ = torch.cat([x2, x2_pad, x2_norm], dim=-1)
    res = x1_.matmul(x2_.transpose(-2, -1))

    # Zero out negative values
    res.clamp_min_(1e-30).sqrt_()
    return res

This uses the same basic strategy as the original method I'd written, but does everything in a single matmul call because matmul calls are really just that fast.

That being said, it has some drawbacks for inclusion upstream in pytorch proper over my original snippet. Namely (1) the additional working memory overhead, and (2) it's the same basic strategy, so a proper implementation of the first snippet I'd linked may erase the gap anyways.

@jacobrgardner Thanks, let me check the performance metrics.

Great work!

FYI, I just added a reference to an implementation for the batched version in #15253 (comment)
The performance metrics were for that version. See facebookresearch/PyTorch-BigGraph#67 (comment).

cc: @lerks

@ifedan, can we also add a squared kwarg (or maybe a torch.sqcdist function) that returns the squared differences? This is important for kernels methods, when we want to differentiate through some function fo the squared distance between the same point (i.e. at zero).

While the squared distance is differentiable at zero, this is not true for the distance. But autograd doesn't know that, so doing cidst(x1, x2).pow(2) is not an option and we need the squared distances directly.

@pytorchbot retest this please

@ifedan any thoughts on being able to return the squared (or more generally p-ed) distance?

That is have an option to return \sum_k (x_ik - xj_k)^p instead of (\sum_k (x_ik - xj_k)^p(^(1/p).

This is not just as easy as squaring the result of cdist, as cdist is not differentiable at zero.

Did you move cdist to functional.py because optional kwargs are not handled well in native_functions?
cc @houseroad for backward compatibility (an optional arg is added to cdist)

@pytorchbot retest this please

@ifedan merged this pull request in 12dde7f.

What do you guys think about https://stackoverflow.com/questions/61241523/replacing-torch-cdist-function-to-eliminate-gpu-out-of-memory-runtime-error ?

def new_cdist(p, eta):
    class cdist(torch.autograd.Function):
        @staticmethod
        def forward(ctx, W, X):
            ctx.save_for_backward(W, X)
            out = -torch.cdist(W, X, p)
            return out

        @staticmethod
        def backward(ctx, grad_output):
            W, X = ctx.saved_tensors
            grad_W = grad_X = None
            if ctx.needs_input_grad[0]:
                _temp1 = torch.unsqueeze(X, 2).expand(X.shape[0], X.shape[1], W.shape[0]).permute(1, 0, 2)
                _temp2 = torch.unsqueeze(W.transpose(0, 1), 1)
                _temp = torch.cdist(_temp1, _temp2, p).squeeze().transpose(0, 1)
                grad_W = torch.matmul(grad_output, _temp)
                # print('before norm: ', torch.norm(grad_W))
                grad_W = eta * np.sqrt(grad_W.numel()) / torch.norm(grad_W) * grad_W
                print('after norm: ', torch.norm(grad_W))
            if ctx.needs_input_grad[1]:
                _temp1 = torch.unsqueeze(W, 2).expand(W.shape[0], W.shape[1], X.shape[0]).permute(1, 0, 2)
                _temp2 = torch.unsqueeze(X.transpose(0, 1), 1)
                _temp = torch.cdist(_temp1, _temp2, p).squeeze().transpose(0, 1)
                _temp = torch.nn.functional.hardtanh(_temp, min_val=-1., max_val=1.)
                grad_X = torch.matmul(grad_output.transpose(0, 1), _temp)
            return grad_W, grad_X
    return cdist().apply

This (or very similar) was implemented in #31167

@ngimel are you saying that I should install the nightly version of pytorch for #31167 ?

@jacobrgardner As for #25799 (comment) , how to replace torch.cat() which contains contiguous tensors ?

x1_ = torch.cat([-2. * x1, x1_norm, x1_pad], dim=-1)
x2_ = torch.cat([x2, x2_pad, x2_norm], dim=-1)

You can install either nightly or 1.5.