LazyTensor

Summary

This section contains the full API documentation of the LazyTensor wrapper, which works identically on NumPy arrays and PyTorch tensors.

Vi(x_or_ind[, dim])	Simple wrapper that returns an instantiation of LazyTensor of type 0.
Vj(x_or_ind[, dim])	Simple wrapper that returns an instantiation of LazyTensor of type 1.
Pm(x_or_ind[, dim])	Simple wrapper that returns an instantiation of LazyTensor of type 2.
LazyTensor([x, axis, is_complex])	Symbolic wrapper for PyTorch tensors.
ComplexLazyTensor([x, axis])	Extension of the LazyTensor class for complex operations.

Syntax

pykeops.torch.Vi(x_or_ind, dim=None)

Simple wrapper that returns an instantiation of LazyTensor of type 0.

pykeops.torch.Vj(x_or_ind, dim=None)

Simple wrapper that returns an instantiation of LazyTensor of type 1.

pykeops.torch.Pm(x_or_ind, dim=None)

Simple wrapper that returns an instantiation of LazyTensor of type 2.

class pykeops.torch.LazyTensor(x=None, axis=None, is_complex=False)

Symbolic wrapper for PyTorch tensors.

LazyTensor encode numerical arrays through the combination of a symbolic, mathematical formula and a list of small data arrays. They can be used to implement efficient algorithms on objects that are easy to define, but impossible to store in memory (e.g. the matrix of pairwise distances between two large point clouds).

LazyTensor may be created from standard NumPy arrays or PyTorch tensors, combined using simple mathematical operations and converted back to NumPy arrays or PyTorch tensors with efficient reduction routines, which outperform standard tensorized implementations by two orders of magnitude.

static __new__(self, x=None, axis=None, is_complex=False)

__init__(x=None, axis=None, is_complex=False)

Creates a KeOps symbolic variable.

Parameters:

• x –

  May be either:

  • A float, a list of floats, a NumPy float, a 0D or 1D NumPy array, a 0D or 1D PyTorch tensor, in which case the LazyTensor represents a constant vector of parameters, to be broadcasted on other LazyTensor.

  • A 2D NumPy array or PyTorch tensor, in which case the LazyTensor represents a variable indexed by \(i\) if axis=0 or \(j\) if axis=1.

  • A 3D+ NumPy array or PyTorch tensor with a dummy dimension (=1) at position -3 or -2, in which case the LazyTensor represents a variable indexed by \(i\) or \(j\), respectively. Dimensions before the last three will be handled as batch dimensions, that may support operator broadcasting.

  • A tuple of 3 integers (ind,dim,cat), in which case the LazyTensor represents a symbolic variable that should be instantiated at call-time.

  • An integer, in which case the LazyTensor represents an integer constant handled efficiently at compilation time.

  • None, for internal use.

• axis (int) – should be equal to 0 or 1 if x is a 2D tensor, and None otherwise.

Warning

A LazyTensor constructed from a NumPy array or a PyTorch tensor retains its dtype (float16, float32 or float64) and device properties (is it stored on the GPU?). Since KeOps does not support automatic type conversions and data transfers, please make sure not to mix LazyTensor that come from different frameworks/devices or which are stored with different precisions.

get_tools()

This method is specialized in pykeops.numpy.LazyTensor and pykeops.torch.LazyTensor. It populates the tools class.

lt_constructor(x=None, axis=None, is_complex=False)

This method is specialized in pykeops.numpy.LazyTensor and pykeops.torch.LazyTensor. It returns a new instance of a LazyTensor (numpy or pytorch).

class pykeops.torch.ComplexLazyTensor(x=None, axis=None)

Extension of the LazyTensor class for complex operations.

__init__(x=None, axis=None)

Creates a KeOps symbolic variable of complex dtype.

get_tools()

This method is specialized in pykeops.numpy.LazyTensor and pykeops.torch.LazyTensor. It populates the tools class.

lt_constructor(x=None, axis=None, is_complex=True)

This method is specialized in pykeops.numpy.LazyTensor and pykeops.torch.LazyTensor. It returns a new instance of a LazyTensor (numpy or pytorch).