Use cdist for distance computations when available.
@@ -35,6 +35,9 @@ class Distance(torch.nn.Module):
# @torch.jit.script_method
def _jit_sq_dist_x1_neq_x2_nobatch(self, x1, x2, postprocess):
+ if hasattr(torch, 'cdist'): # TODO: Remove lower branch when PyTorch 1.1.0 releases.
+ res = torch.cdist(x1, x2).pow(2)
+ else:
# Compute squared distance matrix using quadratic expansion
x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
x2_norm = x2.pow(2).sum(dim=-1, keepdim=True)
@@ -58,6 +61,9 @@ class Distance(torch.nn.Module):
# @torch.jit.script_method
def _jit_sq_dist_x1_eq_x2_nobatch(self, x1, postprocess):
+ if hasattr(torch, 'cdist'): # TODO: Remove lower branch when PyTorch 1.1.0 releases.
+ res = torch.cdist(x1, x1).pow(2)
+ else:
# Compute squared distance matrix using quadratic expansion
x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
res = torch.addmm(x1_norm.transpose(-2, -1), x1, x1.transpose(-2, -1), alpha=-2).add_(x1_norm)
@@ -77,8 +83,9 @@ class Distance(torch.nn.Module):
# @torch.jit.script_method
def _jit_dist_x1_neq_x2_nobatch(self, x1, x2, postprocess, false_tensor):
- # Need to do a hack to get a False tensor because pytorch stable doesn't
- # support creating of tensors in torch scripts
+ if hasattr(torch, 'cdist'):
+ res = torch.cdist(x1, x2)
+ else:
res = self._jit_sq_dist_x1_neq_x2_nobatch(x1, x2, postprocess=false_tensor)
res = res.clamp_min_(1e-30).sqrt_()
return self._postprocess(res) if bool(postprocess) else res
@@ -93,10 +100,12 @@ class Distance(torch.nn.Module):
# @torch.jit.script_method
def _jit_dist_x1_eq_x2_nobatch(self, x1, postprocess, false_tensor):
- # Need to do a hack to get a False tensor because pytorch stable doesn't
- # support creating of tensors in torch scripts
+ if hasattr(torch, 'cdist'):
+ res = torch.cdist(x1, x1)
+ else:
res = self._jit_sq_dist_x1_eq_x2_nobatch(x1, postprocess=false_tensor)
res = res.clamp_min_(1e-30).sqrt_()
+
return self._postprocess(res) if bool(postprocess) else res
