pairwise_distances#

torchdr.pairwise_distances(X: Tensor, Y: Tensor | None = None, metric: str = 'euclidean', backend: str | FaissConfig | None = None, exclude_diag: bool = False, k: int | None = None, return_indices: bool = False, device: str = 'auto')[source]#

Compute pairwise distances between two tensors.

This is the main distance computation function that supports multiple backends for efficient computation. It can compute: - Full pairwise distance matrices between X and Y (or X and itself) - k-nearest neighbor distances when k is specified - Distances with various metrics (euclidean, manhattan, angular, etc.)

For computing distances between specific indexed subsets, use pairwise_distances_indexed instead.

Parameters:

  • X (torch.Tensor of shape (n_samples, n_features)) – Input data.

  • Y (torch.Tensor of shape (m_samples, n_features), optional) – Input data. If None, Y is set to X.

  • metric (str, optional) – Metric to use. Default is “euclidean”.

  • backend ({'keops', 'faiss', None} or FaissConfig, optional) – Backend to use for computation. Can be: - “keops”: Use KeOps for memory-efficient symbolic computations - “faiss”: Use FAISS for fast k-NN computations with default settings - None: Use standard PyTorch operations - FaissConfig object: Use FAISS with custom configuration If None, use standard torch operations.

  • exclude_diag (bool, optional) – Whether to exclude the diagonal from the distance matrix. Only used when k is not None. Default is False.

  • k (int, optional) – If not None, return only the k-nearest neighbors.

  • return_indices (bool, optional) – Whether to return the indices of the k-nearest neighbors. Default is False.

  • device (str, default="auto") – Device to use for computation. If “auto”, keeps data on its current device. Otherwise, temporarily moves data to specified device for computation. Output remains on the computation device. Used with backend=None (torch) and backend=”keops”.

Returns:

  • C (torch.Tensor) – Pairwise distances.

  • indices (torch.Tensor, optional) – Indices of the k-nearest neighbors. Only returned if k is not None.

Examples

>>> import torch
>>> from torchdr.distance import pairwise_distances, FaissConfig

>>> # Basic usage
>>> X = torch.randn(1000, 128)
>>> distances = pairwise_distances(X, k=10, backend='faiss')

>>> # With float16 precision for GPU
>>> config = FaissConfig(use_float16=True)
>>> distances = pairwise_distances(X.cuda(), k=10, backend=config)

>>> # Using FaissConfig with custom settings
>>> config = FaissConfig(use_float16=True, temp_memory=2.0)
>>> distances = pairwise_distances(X.cuda(), k=10, backend=config)
On this page

This Page