LazyTensor

Summary

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

GenericLazyTensor([x, axis])	Symbolic wrapper for NumPy arrays and PyTorch tensors.
ComplexGenericLazyTensor([x, axis])	Extension of the LazyTensor class for complex operations.

Syntax

class pykeops.common.lazy_tensor.GenericLazyTensor(x=None, axis=None)

Symbolic wrapper for NumPy arrays and PyTorch tensors. This is the abstract class, end user should use pykeops.numpy.LazyTensor or pykeops.torch.LazyTensor.

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.

__init__(x=None, axis=None)

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.

lt_constructor(x=None, axis=None)

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

get_tools()

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

fixvariables()

If needed, assigns final labels to each variable and pads their batch dimensions prior to a Genred() call.

promote(other, props, is_complex=False)

Creates a new LazyTensor whose None properties are set to those of self or other.

init(is_complex=False)

Creates a copy of a LazyTensor, without formula attribute.

join(other, is_complex=False)

Merges the variables and attributes of two LazyTensor, with a compatibility check. This method concatenates tuples of variables, without paying attention to repetitions.

unary(operation, dimres=None, opt_arg=None, opt_arg2=None, is_complex=None)

Symbolically applies operation to self, with optional arguments if needed.

The optional argument dimres may be used to specify the dimension of the output result.

binary(other, operation, is_operator=False, dimres=None, dimcheck='sameor1', opt_arg=None, opt_pos='last', rversion=False, is_complex=None)

Symbolically applies operation to self, with optional arguments if needed.

Keyword Arguments:

• dimres (-) – May be used to specify the dimension of the output result.

• is_operator (-) – May be used to specify if operation is an operator like +, - or a “genuine” function.

• dimcheck (-) – shall we check the input dimensions? Supported values are "same", "sameor1", "vecand1" or None.

• rversion (-) – shall we invert lhs and rhs of the binary op, e.g. as in __radd__, __rmut__, etc…

ternary(other1, other2, operation, dimres=None, dimcheck='sameor1', opt_arg=None)

Symbolically applies operation to self, with optional arguments if needed.

Keyword Arguments:

• dimres (-) – May be used to specify the dimension of the output result.

• is_operator (-) – May be used to specify if operation is an operator like +, - or a “genuine” function.

• dimcheck (-) – shall we check the input dimensions? Supported values are "same", "sameor1", or None.

reduction(reduction_op, other=None, opt_arg=None, axis=None, dim=None, call=True, is_complex=None, **kwargs)

Applies a reduction to a LazyTensor. This method is used internally by the LazyTensor class. :param reduction_op: the string identifier of the reduction, which will be passed to the KeOps routines. :type reduction_op: string

Keyword Arguments:

• other – May be used to specify some weights ; depends on the reduction.

• opt_arg – typically, some integer needed by ArgKMin reductions ; depends on the reduction.

• axis (integer) – The axis with respect to which the reduction should be performed. Supported values are nbatchdims and nbatchdims + 1, where nbatchdims is the number of “batch” dimensions before the last three (\(i\) indices, \(j\) indices, variables’ dimensions).

• dim (integer) – alternative keyword for the axis argument.

• call (True or False) – Should we actually perform the reduction on the current variables? If True, the returned object will be a NumPy array or a PyTorch tensor. Otherwise, we simply return a callable LazyTensor that may be used as a pykeops.numpy.Genred or pykeops.torch.Genred function on arbitrary tensor data.

• backend (string) – Specifies the map-reduce scheme, as detailed in the documentation of the Genred module.

• device_id (int, default=-1) – Specifies the GPU that should be used to perform the computation; a negative value lets your system choose the default GPU. This parameter is only useful if your system has access to several GPUs.

• ranges (6-uple of IntTensors, None by default) – Ranges of integers that specify a block-sparse reduction scheme as detailed in the documentation of the Genred module. If None (default), we simply use a dense Kernel matrix as we loop over all indices \(i\in[0,M)\) and \(j\in[0,N)\).

• dtype_acc (string, default "auto") –

  type for accumulator of reduction, before casting to dtype. It improves the accuracy of results in case of large sized data, but is slower. Default value “auto” will set this option to the value of dtype. The supported values are:

  • dtype_acc = "float16" : allowed only if dtype is “float16”.

  • dtype_acc = "float32" : allowed only if dtype is “float16” or “float32”.

  • dtype_acc = "float64" : allowed only if dtype is “float32” or “float64”.

• use_double_acc (bool, default False) – same as setting dtype_acc=”float64” (only one of the two options can be set). If True, accumulate results of reduction in float64 variables, before casting to float32. This can only be set to True when data is in float32 or float64. It improves the accuracy of results in case of large sized data, but is slower.

• sum_scheme (string, default "auto") –

  method used to sum up results for reductions. This option may be changed only when reduction_op is one of: “Sum”, “MaxSumShiftExp”, “LogSumExp”, “Max_SumShiftExpWeight”, “LogSumExpWeight”, “SumSoftMaxWeight”. Default value “auto” will set this option to “block_red” for these reductions. Possible values are:

  • sum_scheme = "direct_sum": direct summation

  • sum_scheme = "block_sum": use an intermediate accumulator in each block before accumulating in the output. This improves accuracy for large sized data.

  • sum_scheme = "kahan_scheme": use Kahan summation algorithm to compensate for round-off errors. This improves accuracy for large sized data.

• enable_chunks (bool, default True) – enable automatic selection of special “chunked” computation mode for accelerating reductions with formulas involving large dimension variables.

• out (2d NumPy array or PyTorch Tensor, None by default) – The output numerical array, for in-place computation. If provided, the output array should all have the same dtype, be contiguous and be stored on the same device as the arguments. Moreover it should have the correct shape for the output.

solve(other, var=None, call=True, **kwargs)

Solves a positive definite linear system of the form sum(self) = other or sum(self*var) = other , using a conjugate gradient solver.

Parameters:

• self (LazyTensor) – KeOps variable that encodes a symmetric positive definite matrix / linear operator.

• other (LazyTensor) – KeOps variable that encodes the second member of the equation.

Keyword Arguments:

• var (LazyTensor) – If var is None, solve will return the solution of the self * var = other equation. Otherwise, if var is a KeOps symbolic variable, solve will assume that self defines an expression that is linear with respect to var and solve the equation self(var) = other with respect to var.

• alpha (float, default=1e-10) – Non-negative ridge regularization parameter.

• call (bool) – If True and if no other symbolic variable than var is contained in self, solve will return a tensor solution of our linear system. Otherwise solve will return a callable LazyTensor.

• backend (string) – Specifies the map-reduce scheme, as detailed in the documentation of the Genred module.

• device_id (int, default=-1) – Specifies the GPU that should be used to perform the computation; a negative value lets your system choose the default GPU. This parameter is only useful if your system has access to several GPUs.

• ranges (6-uple of IntTensors, None by default) – Ranges of integers that specify a block-sparse reduction scheme as detailed in the documentation of the Genred module. If None (default), we simply use a dense Kernel matrix as we loop over all indices \(i\in[0,M)\) and \(j\in[0,N)\).

• dtype_acc (string, default "auto") –

  type for accumulator of reduction, before casting to dtype. It improves the accuracy of results in case of large sized data, but is slower. Default value “auto” will set this option to the value of dtype. The supported values are:

  • dtype_acc = "float16" : allowed only if dtype is “float16”.

  • dtype_acc = "float32" : allowed only if dtype is “float16” or “float32”.

  • dtype_acc = "float64" : allowed only if dtype is “float32” or “float64”.

• use_double_acc (bool, default False) – same as setting dtype_acc=”float64” (only one of the two options can be set) If True, accumulate results of reduction in float64 variables, before casting to float32. This can only be set to True when data is in float32 or float64. It improves the accuracy of results in case of large sized data, but is slower.

• sum_scheme (string, default "auto") –

  method used to sum up results for reductions. This option may be changed only when reduction_op is one of: “Sum”, “MaxSumShiftExp”, “LogSumExp”, “Max_SumShiftExpWeight”, “LogSumExpWeight”, “SumSoftMaxWeight”. Default value “auto” will set this option to “block_red” for these reductions. Possible values are:

  • sum_scheme = "direct_sum": direct summation

  • sum_scheme = "block_sum": use an intermediate accumulator in each block before accumulating in the output. This improves accuracy for large sized data.

  • sum_scheme = "kahan_scheme": use Kahan summation algorithm to compensate for round-off errors. This improves accuracy for large sized data.

• enable_chunks (bool, default True) – enable automatic selection of special “chunked” computation mode for accelerating reductions with formulas involving large dimension variables.

Warning

Please note that no check of symmetry and definiteness will be performed prior to our conjugate gradient descent.

__call__(*args, **kwargs)

Executes a Genred or KernelSolve call on the input data, as specified by self.formula .

__str__()

Returns a verbose string identifier.

property shape

Returns the shape of the LazyTensor

dim()

Just as in PyTorch, returns the number of dimensions of a LazyTensor.

__add__(other)

Broadcasted addition operator - a binary operation.

x + y returns a LazyTensor that encodes, symbolically, the addition of x and y.

__radd__(other)

Broadcasted addition operator - a binary operation.

x + y returns a LazyTensor that encodes, symbolically, the addition of x and y.

__sub__(other)

Broadcasted subtraction operator - a binary operation.

x - y returns a LazyTensor that encodes, symbolically, the subtraction of x and y.

__rsub__(other)

Broadcasted subtraction operator - a binary operation.

x - y returns a LazyTensor that encodes, symbolically, the subtraction of x and y.

__mul__(other)

Broadcasted element-wise product - a binary operation.

x * y returns a LazyTensor that encodes, symbolically, the elementwise product of x and y.

__rmul__(other)

Broadcasted element-wise product - a binary operation.

x * y returns a LazyTensor that encodes, symbolically, the elementwise product of x and y.

__truediv__(other)

Broadcasted element-wise division - a binary operation.

x / y returns a LazyTensor that encodes, symbolically, the elementwise division of x by y.

__rtruediv__(other)

Broadcasted element-wise division - a binary operation.

x / y returns a LazyTensor that encodes, symbolically, the elementwise division of x by y.

__or__(other)

Euclidean scalar product - a binary operation.

(x|y) returns a LazyTensor that encodes, symbolically, the scalar product of x and y which are assumed to have the same shape.

__ror__(other)

Euclidean scalar product - a binary operation.

(x|y) returns a LazyTensor that encodes, symbolically, the scalar product of x and y which are assumed to have the same shape.

__abs__()

Element-wise absolute value - a unary operation.

abs(x) returns a LazyTensor that encodes, symbolically, the element-wise absolute value of x.

abs()

Element-wise absolute value - a unary operation.

x.abs() returns a LazyTensor that encodes, symbolically, the element-wise absolute value of x.

__neg__()

Element-wise minus - a unary operation.

-x returns a LazyTensor that encodes, symbolically, the element-wise opposite of x.

exp()

Element-wise exponential - a unary operation.

x.exp() returns a LazyTensor that encodes, symbolically, the element-wise exponential of x.

log()

Element-wise logarithm - a unary operation.

x.log() returns a LazyTensor that encodes, symbolically, the element-wise logarithm of x.

xlogx()

Element-wise x*log(x) function - a unary operation.

x.xlogx() returns a LazyTensor that encodes, symbolically, the element-wise x times logarithm of x (with value 0 at 0).

cos()

Element-wise cosine - a unary operation.

x.cos() returns a LazyTensor that encodes, symbolically, the element-wise cosine of x.

sin()

Element-wise sine - a unary operation.

x.sin() returns a LazyTensor that encodes, symbolically, the element-wise sine of x.

sinxdivx()

Element-wise sin(x)/x function - a unary operation.

x.sinxdivx() returns a LazyTensor that encodes, symbolically, the element-wise sinxdivx function of x.

sinc()

Element-wise sinc(x) = sin(pi x) / (pi x) function - a unary operation.

x.sinc() returns a LazyTensor that encodes, symbolically, the element-wise sinc function of x.

asin()

Element-wise arcsine - a unary operation.

x.asin() returns a LazyTensor that encodes, symbolically, the element-wise arcsine of x.

acos()

Element-wise arccosine - a unary operation.

x.acos() returns a LazyTensor that encodes, symbolically, the element-wise arccosine of x.

atan()

Element-wise arctangent - a unary operation.

x.atan() returns a LazyTensor that encodes, symbolically, the element-wise arctangent of x.

atan2(other)

Element-wise atan2 - a binary operation.

y.atan2(x) returns a LazyTensor that encodes, symbolically, the element-wise atan2 of x and y.

sqrt()

Element-wise square root - a unary operation.

x.sqrt() returns a LazyTensor that encodes, symbolically, the element-wise square root of x.

rsqrt()

Element-wise inverse square root - a unary operation.

x.rsqrt() returns a LazyTensor that encodes, symbolically, the element-wise inverse square root of x.

__pow__(other)

Broadcasted element-wise power operator - a binary operation.

x**y returns a LazyTensor that encodes, symbolically, the element-wise value of x to the power y.

Note

• if y = 2, x**y relies on the "Square" KeOps operation;

• if y = 0.5, x**y uses on the "Sqrt" KeOps operation;

• if y = -0.5, x**y uses on the "Rsqrt" KeOps operation.

power(other)

Broadcasted element-wise power operator - a binary operation.

pow(x,y) is equivalent to x**y.

square()

Element-wise square - a unary operation.

x.square() is equivalent to x**2 and returns a LazyTensor that encodes, symbolically, the element-wise square of x.

__eq__(other)

“equal to” - a binary operation.

x==y returns a LazyTensor that encodes, symbolically, the element-wise operation x==y (with outputs values 0 or 1).

__ne__(other)

“not equal to” - a binary operation.

x!=y returns a LazyTensor that encodes, symbolically, the element-wise operation x!=y (with outputs values 0 or 1).

__lt__(other)

“less than” - a binary operation.

x<y returns a LazyTensor that encodes, symbolically, the element-wise operation x<y (with outputs values 0 or 1).

__gt__(other)

“greater than” - a binary operation.

x>y returns a LazyTensor that encodes, symbolically, the element-wise operation x>y (with outputs values 0 or 1).

__le__(other)

“less than or equal to” - a binary operation.

x<=y returns a LazyTensor that encodes, symbolically, the element-wise operation x<=y (with outputs values 0 or 1).

__ge__(other)

“greater than” - a binary operation.

x>=y returns a LazyTensor that encodes, symbolically, the element-wise operation x>=y (with outputs values 0 or 1).

sign()

Element-wise sign in {-1,0,+1} - a unary operation.

x.sign() returns a LazyTensor that encodes, symbolically, the element-wise sign of x.

step()

Element-wise step function - a unary operation.

x.step() returns a LazyTensor that encodes, symbolically, the element-wise sign of x.

relu()

Element-wise ReLU function - a unary operation.

x.relu() returns a LazyTensor that encodes, symbolically, the element-wise positive part of x.

clamp(other1, other2)

Element-wise Clamp function - a ternary operation.

x.clamp(a,b) returns a LazyTensor that encodes, symbolically, the element-wise clamping of x in (a,b). Broadcasting rules apply. a and b may be fixed integers or floats, or other LazyTensors.

ifelse(other1, other2)

Element-wise if-else function - a ternary operation.

x.ifelse(a,b) returns a LazyTensor that encodes, symbolically, a where x >= 0 and b where x < 0. Broadcasting rules apply. a and b may be fixed integers or floats, or other LazyTensors.

mod(modulus, offset=0)

Element-wise modulo with offset function - a ternary operation.

x.mod(a,b) returns a LazyTensor that encodes, symbolically, the element-wise modulo of x with modulus a and offset b. By default b=0, so that x.mod(a) becomes equivalent to the NumPy function mod. Broadcasting rules apply. a and b are fixed integers or float.

round(other=0)

Element-wise rounding function - a unary operation.

x.round(d) returns a LazyTensor that encodes, symbolically, the element-wise rounding of x to d decimal places. d is int.

sqnorm2()

Squared Euclidean norm - a unary operation.

x.sqnorm2() returns a LazyTensor that encodes, symbolically, the squared Euclidean norm of a vector x.

norm2()

Euclidean norm - a unary operation.

x.norm2() returns a LazyTensor that encodes, symbolically, the Euclidean norm of a vector x.

norm(dim)

Euclidean norm - a unary operation.

x.norm(-1) is equivalent to x.norm2() and returns a LazyTensor that encodes, symbolically, the Euclidean norm of a vector x.

normalize()

Vector normalization - a unary operation.

x.normalize() returns a LazyTensor that encodes, symbolically, a vector x divided by its Euclidean norm.

sqdist(other)

Squared distance - a binary operation.

x.sqdist(y) returns a LazyTensor that encodes, symbolically, the squared Euclidean distance between two vectors x and y.

weightedsqnorm(other)

Weighted squared norm of a LazyTensor x - a binary operation.

x.weightedsqnorm(s) returns a LazyTensor that encodes, symbolically, the weighted squared Norm of a vector x with weights stored in the LazyTensor s- see the main reference page for details.

weightedsqdist(g, s)

Weighted squared distance.

x.weightedsqdist(y, s) is equivalent to (x - y).weightedsqnorm(s).

softdtw_sqdist(other, gamma)

Soft-DTW distance (based on the squared euclidean distance)

x.softdtw_sqdist(y, gamma) computes the Soft-DTW distance between vectors x and y` with smoothness parameter gamma, the base dissimilarity being given by the sqaured euclidean distance.

elem(i)

Indexing of a vector - a unary operation.

x.elem(i) returns a LazyTensor that encodes, symbolically, the i-th element x[i] of the vector x.

extract(i, d)

Range indexing - a unary operation.

x.extract(i, d) returns a LazyTensor that encodes, symbolically, the sub-vector x[i:i+d] of the vector x.

__getitem__(key)

Element or range indexing - a unary operation.

x[key] redirects to the elem() or extract() methods, depending on the key argument. Supported values are:

• an integer k, in which case x[key] redirects to elem(x,k),

• a tuple ..,:,:,k with k an integer, which is equivalent to the case above,

• a slice of the form k:l, k: or :l, with k and l two integers, in which case x[key] redirects to extract(x,k,l-k),

• a tuple of slices of the form ..,:,:,k:l, ..,:,:,k: or ..,:,:,:l, with k and l two integers, which are equivalent to the case above.

one_hot(D)

Encodes a (rounded) scalar value as a one-hot vector of dimension D.

x.one_hot(D) returns a LazyTensor that encodes, symbolically, a vector of length D whose round(x)-th coordinate is equal to 1, and the other ones to zero.

bspline(x, k=0)

Vector of BSpline functions of order k for the knots (self), evaluated at x.

Parameters:

• x – a LazyTensor of dimension 1.

• k – a non-negative integer.

concat(other)

Concatenation of two LazyTensor - a binary operation.

x.concat(y) returns a LazyTensor that encodes, symbolically, the concatenation of x and y along their last dimension.

static concatenate(tuple_of_lt, axis=-1)

Concatenation of a tuple of GenericLazyTensor.

GenericLazyTensor.concatenate( (x_1, x_2, ..., x_n), -1) returns a GenericLazyTensor that encodes, symbolically, the concatenation of x_1, x_2, …, x_n along their last dimension. Note that axis should be equal to -1 or 2 (if the x_i’s are 3D GenericLazyTensor): GenericLazyTensors only support concatenation and indexing operations with respect to the last dimension.

static cat(tuple_of_lt, dim)

Concatenation of a tuple of LazyTensors.

LazyTensor.cat( (x_1, x_2, ..., x_n), -1) is a PyTorch-friendly alias for LazyTensor.concatenate( (x_1, x_2, ..., x_n), -1); just like indexing operations, it is only supported along the last dimension.

matvecmult(other)

Matrix-vector product - a binary operation.

If x._shape[-1] == A*B and y._shape[-1] == B, z = x.matvecmult(y) returns a GenericLazyTensor such that z._shape[-1] == A which encodes, symbolically, the matrix-vector product of x and y along their last dimension. For details, please check the documentation of the KeOps operation "MatVecMult" in the main reference page.

vecmatmult(other)

Vector-matrix product - a binary operation.

If x._shape[-1] == A and y._shape[-1] == A*B, z = x.vecmatmult(y) returns a GenericLazyTensor such that z._shape[-1] == B which encodes, symbolically, the vector-matrix product of x and y along their last dimension. For details, please check the documentation of the KeOps operation "VecMacMult" in the main reference page.

tensorprod(other)

Tensor product of vectors - a binary operation.

If x._shape[-1] == A and y._shape[-1] == B, z = x.tensorprod(y) returns a GenericLazyTensor such that z._shape[-1] == A*B which encodes, symbolically, the tensor product of x and y along their last dimension. For details, please check the documentation of the KeOps operation "TensorProd" in the main reference page.

keops_tensordot(other, dimfa, dimfb, contfa, contfb, *args)

Tensor dot product (on KeOps internal dimensions) - a binary operation.

Parameters:

• other – a LazyTensor

• dimfa – tuple/list of int

• dimfb – tuple/list of int

• contfa – tuple/list of int listing contraction dimension of a (could be empty)

• contfb – tuple/list of int listing contraction dimension of b (could be empty)

• args – a tuple/list of int containing the graph of a permutation of the output

Returns:

a LazyTensor

keops_kron(other, dimfself, dimfother)

Kronecker product (on KeOps internal dimensions) - a binary operation.

If self._shape[-1] == d0 * d1 * ... * dN and other._shape[-1] == D0 * D1 * ... * DN, z = x.keops_kron(y, [d0, d1, ..., dN], [D0, D1, ..., DN]) returns a GenericLazyTensor of shape z._shape[-1] == d0 * D0 * d1 * D1 * ... * dN * DN which encodes, symbolically, the (flattened version of) Kronecker product of self and other along their internal dimension.

Parameters:

• other – a LazyTensor

• dimfself – tuple/list of int listing the dimensions of a (prod(dimfself) == self._shape[-1]) and len(dimfself) == len(dimfother)

• dimfother – tuple/list of int listing the dimensions of b (prod(dimfother) == y._shape[-1]) and len(dimfself) == len(dimfother)

Returns:

a LazyTensor of internal shape prod(dimfself) * prod(dimfother)

>>>  import numpy as np
>>>  from pykeops.numpy import LazyTensor
>>>
>>>  dimfb = [3, 3, 2, 1]
>>>  dimfa = [3, 1, 3, 2]
>>>
>>>  M, N, axis = 11, 150, 1
>>>
>>>  x = np.random.rand(M, np.array(dimfa).prod())
>>>  y = np.random.rand(N, np.array(dimfb).prod())
>>>
>>>  gamma_py = np.zeros((M, N, np.array(dimfa).prod() * np.array(dimfb).prod()))
>>>  for i in range(M):
>>>      for j in range(N):
>>>          gamma_py[i, j, :] = np.kron(x[i, :].reshape(*dimfa), y[j, :].reshape(*dimfb)).flatten()
>>>
>>>  gamma_py = gamma_py.sum(axis=axis)
>>>
>>>  X = LazyTensor(x[:, None, :])
>>>  Y = LazyTensor(y[None, :, :])
>>>
>>>  gamma_keops = (X.keops_kron(Y, dimfa, dimfb)).sum(axis=axis)
>>>
>>>  print(np.allclose(gamma_keops, gamma_py, atol=1e-6))

grad(other, gradin)

Symbolic gradient operation.

z = x.grad(v,e) returns a LazyTensor which encodes, symbolically, the gradient (more precisely, the adjoint of the differential operator) of x, with respect to variable v, and applied to e. For details, please check the documentation of the KeOps operation "Grad" in the main reference page.

diff(other, diffin)

Symbolic differential operation (forward differentiation).

z = x.diff(v,e) returns a LazyTensor which encodes, symbolically, the differential operator of x, with respect to variable v, and applied to e. For details, please check the documentation of the KeOps operation "Diff" in the main reference page.

factorize(other)

Symbolic factorization operation.

z = x.factorize(y) returns a LazyTensor which encodes, symbolically, a factorization of x with respect to y. For details, please check the documentation of the KeOps operation "Factorize" in the main reference page.

auto_factorize()

Symbolic auto factorization operation.

z = x.auto_factorize() returns a LazyTensor which encodes, symbolically, the automatic factorization of x. For details, please check the documentation of the KeOps operation "AutoFactorize" in the main reference page.

grad_matrix(other)

Symbolic gradient matrix operation.

z = x.grad_matrix(v) returns a LazyTensor which encodes, symbolically, the gradient matrix (more precisely, the adjoint of the differential operator) of x, with respect to variable v. For details, please check the documentation of the KeOps operation "GradMatrix" in the main reference page.

trace_operator(var)

Symbolic trace operation.

z = x.trace_operator(v) returns a LazyTensor which encodes, symbolically, the trace of x, with respect to variable v. x must be a linear function of v. For details, please check the documentation of the KeOps operation "TraceOperator" in the main reference page.

divergence(var)

Symbolic divergence operation.

z = x.divergence(v) returns a LazyTensor which encodes, symbolically, the divergence of x, with respect to variable v. Inner dimensions of x (x.ndim) and v (v.ndim) must match.

laplacian(var)

Symbolic laplacian operation.

z = x.laplacian(v) returns a LazyTensor which encodes, symbolically, the laplacian of x, with respect to variable v. Inner dimension of x (x.ndim) must equal 1.

sum(axis=-1, dim=None, **kwargs)

Summation unary operation, or Sum reduction.

sum(axis, dim, **kwargs) will:

• if axis or dim = 0, return the sum reduction of self over the “i” indexes.

• if axis or dim = 1, return the sum reduction of self over the “j” indexes.

• if axis or dim = 2, return a new LazyTensor object representing the sum of the values of the vector self,

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), 1 (= reduction over \(j\)) or 2 (i.e. -1, sum along the dimension of the vector variable).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

sum_reduction(axis=None, dim=None, **kwargs)

Sum reduction.

sum_reduction(axis, dim, **kwargs) will return the sum reduction of self.

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

logsumexp(axis=None, dim=None, weight=None, **kwargs)

Log-Sum-Exp reduction.

logsumexp(axis, dim, weight, **kwargs) will:

• if axis or dim = 0, return the “log-sum-exp” reduction of self over the “i” indexes.

• if axis or dim = 1, return the “log-sum-exp” reduction of self over the “j” indexes.

For details, please check the documentation of the KeOps reductions LogSumExp and LogSumExpWeight in the main reference page.

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• weight (LazyTensor) – optional object that specifies scalar or vector-valued weights in the log-sum-exp operation

• **kwargs – optional parameters that are passed to the reduction() method.

logsumexp_reduction(**kwargs)

Log-Sum-Exp reduction. Redirects to logsumexp() method.

sumsoftmaxweight(weight, axis=None, dim=None, **kwargs)

Sum of weighted Soft-Max reduction.

sumsoftmaxweight(weight, axis, dim, **kwargs) will:

• if axis or dim = 0, return the “sum of weighted Soft-Max” reduction of self over the “i” indexes.

• if axis or dim = 1, return the “sum of weighted Soft-Max” reduction of self over the “j” indexes.

For details, please check the documentation of the KeOps reduction SumSoftMaxWeight in the main reference page.

Keyword Arguments:

• weight (LazyTensor) – object that specifies scalar or vector-valued weights.

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

sumsoftmaxweight_reduction(**kwargs)

Sum of weighted Soft-Max reduction. Redirects to sumsoftmaxweight() method.

min(axis=-1, dim=None, **kwargs)

Minimum unary operation, or Min reduction.

min(axis, dim, **kwargs) will:

• if axis or dim = 0, return the min reduction of self over the “i” indexes.

• if axis or dim = 1, return the min reduction of self over the “j” indexes.

• if axis or dim = 2, return a new LazyTensor object representing the min of the values of the vector self,

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), 1 (= reduction over \(j\)) or 2 (i.e. -1, min along the dimension of the vector variable).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

min_reduction(axis=None, dim=None, **kwargs)

Min reduction.

min_reduction(axis, dim, **kwargs) will return the min reduction of self.

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

__min__(**kwargs)

Minimum unary operation, or Min reduction. Redirects to min() method.

argmin(axis=-1, dim=None, **kwargs)

ArgMin unary operation, or ArgMin reduction.

argmin(axis, dim, **kwargs) will:

• if axis or dim = 0, return the argmin reduction of self over the “i” indexes.

• if axis or dim = 1, return the argmin reduction of self over the “j” indexes.

• if axis or dim = 2, return a new LazyTensor object representing the argmin of the values of the vector self,

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), 1 (= reduction over \(j\)) or 2 (i.e. -1, argmin along the dimension of the vector variable).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

argmin_reduction(axis=None, dim=None, **kwargs)

ArgMin reduction.

argmin_reduction(axis, dim, **kwargs) will return the argmin reduction of self.

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

min_argmin(axis=None, dim=None, **kwargs)

Min-ArgMin reduction.

min_argmin(axis, dim, **kwargs) will:

• if axis or dim = 0, return the minimal values and its indices of self over the “i” indexes.

• if axis or dim = 1, return the minimal values and its indices of self over the “j” indexes.

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

min_argmin_reduction(**kwargs)

Min-ArgMin reduction. Redirects to min_argmin() method.

max(axis=-1, dim=None, **kwargs)

Miaximum unary operation, or Max reduction.

max(axis, dim, **kwargs) will:

• if axis or dim = 0, return the max reduction of self over the “i” indexes.

• if axis or dim = 1, return the max reduction of self over the “j” indexes.

• if axis or dim = 2, return a new LazyTensor object representing the max of the values of the vector self,

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), 1 (= reduction over \(j\)) or 2 (i.e. -1, max along the dimension of the vector variable).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

max_reduction(axis=None, dim=None, **kwargs)

Max reduction.

max_reduction(axis, dim, **kwargs) will return the max reduction of self.

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

__max__(**kwargs)

Maximum unary operation, or Max reduction. Redirects to max() method.

argmax(axis=-1, dim=None, **kwargs)

ArgMax unary operation, or ArgMax reduction.

argmax(axis, dim, **kwargs) will:

• if axis or dim = 0, return the argmax reduction of self over the “i” indexes.

• if axis or dim = 1, return the argmax reduction of self over the “j” indexes.

• if axis or dim = 2, return a new LazyTensor object representing the argmax of the values of the vector self,

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), 1 (= reduction over \(j\)) or 2 (i.e. -1, argmax along the dimension of the vector variable).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

argmax_reduction(axis=None, dim=None, **kwargs)

ArgMax reduction.

argmax_reduction(axis, dim, **kwargs) will return the argmax reduction of self.

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

max_argmax(axis=None, dim=None, **kwargs)

Max-ArgMax reduction.

max_argmax(axis, dim, **kwargs) will:

• if axis or dim = 0, return the maximal values and its indices of self over the “i” indexes.

• if axis or dim = 1, return the maximal values and its indices of self over the “j” indexes.

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

max_argmax_reduction(**kwargs)

Max-ArgMax reduction. Redirects to max_argmax() method.

Kmin(K, axis=None, dim=None, **kwargs)

K-Min reduction.

Kmin(K, axis, dim, **kwargs) will:

• if axis or dim = 0, return the K minimal values of self over the “i” indexes.

• if axis or dim = 1, return the K minimal values of self over the “j” indexes.

Keyword Arguments:

• K (integer) – number of minimal values required

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

Kmin_reduction(**kwargs)

Kmin reduction. Redirects to Kmin() method.

argKmin(K, axis=None, dim=None, **kwargs)

argKmin reduction.

argKmin(K, axis, dim, **kwargs) will:

• if axis or dim = 0, return the indices of the K minimal values of self over the “i” indexes.

• if axis or dim = 1, return the indices of the K minimal values of self over the “j” indexes.

Keyword Arguments:

• K (integer) – number of minimal values required

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

argKmin_reduction(**kwargs)

argKmin reduction. Redirects to argKmin() method.

Kmin_argKmin(K, axis=None, dim=None, **kwargs)

K-Min-argK-min reduction.

Kmin_argKmin(K, axis, dim, **kwargs) will:

• if axis or dim = 0, return the K minimal values and its indices of self over the “i” indexes.

• if axis or dim = 1, return the K minimal values and its indices of self over the “j” indexes.

Keyword Arguments:

• K (integer) – number of minimal values required

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), or 1 (= reduction over \(j\)).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

Kmin_argKmin_reduction(**kwargs)

Kmin_argKmin reduction. Redirects to Kmin_argKmin() method.

__matmul__(v, **kwargs)

Matrix-vector or Matrix-matrix product, supporting batch dimensions.

If K is a LazyTensor whose trailing dimension K._shape[-1] is equal to 1, we can understand it as a linear operator and apply it to arbitrary NumPy arrays or PyTorch Tensors. Assuming that v is a 1D (resp. ND) tensor such that K.shape[-1] == v.shape[-1] (resp. v.shape[-2]), K @ v denotes the matrix-vector (resp. matrix-matrix) product between the two objects, encoded as a vanilla NumPy or PyTorch 1D (resp. ND) tensor.

Example

>>> x, y = torch.randn(1000, 3), torch.randn(2000, 3)
>>> x_i, y_j = LazyTensor( x[:,None,:] ), LazyTensor( y[None,:,:] )
>>> K = (- ((x_i - y_j)**2).sum(2) ).exp()  # Symbolic (1000,2000,1) Gaussian kernel matrix
>>> v = torch.rand(2000, 2)
>>> print( (K @ v).shape )
... torch.Size([1000, 2])

t()

Matrix transposition, permuting the axes of \(i\)- and \(j\)-variables.

For instance, if K is a LazyTensor of shape (B,M,N,D), K.t() returns a symbolic copy of K whose axes 1 and 2 have been switched with each other: K.t().shape == (B,N,M,D).

Example

>>> x, y = torch.randn(1000, 3), torch.randn(2000, 3)
>>> x_i, y_j = LazyTensor( x[:,None,:] ), LazyTensor( y[None,:,:] )
>>> K  = (- ((    x_i     -      y_j   )**2).sum(2) ).exp()  # Symbolic (1000,2000) Gaussian kernel matrix
>>> K_ = (- ((x[:,None,:] - y[None,:,:])**2).sum(2) ).exp()  # Explicit (1000,2000) Gaussian kernel matrix
>>> w  = torch.rand(1000, 2)
>>> print( (K.t() @ w - K_.t() @ w).abs().mean() )
... tensor(1.7185e-05)

property T

Numpy-friendly alias for the matrix transpose self.t().

matvec(v)

Alias for the matrix-vector product, added for compatibility with scipy.sparse.linalg.

If K is a LazyTensor whose trailing dimension K._shape[-1] is equal to 1, we can understand it as a linear operator and wrap it into a scipy.sparse.linalg.LinearOperator object, thus getting access to robust solvers and spectral routines.

Example

>>> import numpy as np
>>> x = np.random.randn(1000,3)
>>> x_i, x_j = LazyTensor( x[:,None,:] ), LazyTensor( x[None,:,:] )
>>> K_xx = (- ((x_i - x_j)**2).sum(2) ).exp()  # Symbolic (1000,1000) Gaussian kernel matrix
>>> from scipy.sparse.linalg import eigsh, aslinearoperator
>>> eigenvalues, eigenvectors = eigsh( aslinearoperator( K_xx ), k=5 )
>>> print(eigenvalues)
... [ 35.5074527   59.01096445  61.35075268  69.34038814 123.77540277]
>>> print( eigenvectors.shape)
... (1000, 5)

rmatvec(v)

Alias for the transposed matrix-vector product, added for compatibility with scipy.sparse.linalg.

See matvec() for further reference.

__hash__ = None

__weakref__

list of weak references to the object

real2complex()

Element-wise “real 2 complex” operation - a unary operation.

x.real2complex() returns a ComplexLazyTensor that encodes, symbolically, the same tensor as x, but seen as complex-valued (with zero imaginary part for each coefficient)

imag2complex()

Element-wise “imag 2 complex” operation - a unary operation.

x.real2complex() returns a ComplexLazyTensor that encodes, symbolically, the multiplication of 1j with x.

exp1j()

Element-wise “complex exponential of 1j x” operation - a unary operation.

x.exp1j() returns a ComplexLazyTensor that encodes, symbolically, the complex exponential of 1j*x.

class pykeops.common.lazy_tensor.ComplexGenericLazyTensor(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.

property shape

returns the shape of the complex LazyTensor.

property real

Element-wise real part of complex - a unary operation.

z.real returns a LazyTensor that encodes, symbolically, the element-wise real part of z.

property imag

Element-wise imaginary part of complex - a unary operation.

z.imag returns a LazyTensor that encodes, symbolically, the element-wise imaginary part of z.

angle()

Element-wise angle (or argument) of complex - a unary operation.

z.angle() returns a LazyTensor that encodes, symbolically, the element-wise angle of z.

conj()

Element-wise complex conjugate - a unary operation.

z.conj() returns a ComplexLazyTensor that encodes, symbolically, the element-wise complex conjugate of z.

sum(axis=-1, dim=None, **kwargs)

Summation unary operation, or Sum reduction.

sum(axis, dim, **kwargs) will:

• if axis or dim = 0, return the sum reduction of self over the “i” indexes.

• if axis or dim = 1, return the sum reduction of self over the “j” indexes.

• if axis or dim = 2, return a new LazyTensor object representing the sum of the values of the vector self,

Keyword Arguments:

• axis (integer) – reduction dimension, which should be equal to the number of batch dimensions plus 0 (= reduction over \(i\)), 1 (= reduction over \(j\)) or 2 (i.e. -1, sum along the dimension of the vector variable).

• dim (integer) – alternative keyword for the axis parameter.

• **kwargs – optional parameters that are passed to the reduction() method.

__abs__()

Element-wise absolute value (or modulus) of complex - a unary operation.

z.abs() returns a LazyTensor that encodes, symbolically, the element-wise absolute value of z.

exp()

Element-wise complex exponential - a unary operation.

z.exp() returns a ComplexLazyTensor that encodes, symbolically, the element-wise complex exponential of z.

real2complex()

Element-wise “real 2 complex” operation - a unary operation.

x.real2complex() returns a ComplexLazyTensor that encodes, symbolically, the same tensor as x, but seen as complex-valued (with zero imaginary part for each coefficient)

imag2complex()

Element-wise “imag 2 complex” operation - a unary operation.

x.real2complex() returns a ComplexLazyTensor that encodes, symbolically, the multiplication of 1j with x.

exp1j()

Element-wise “complex exponential of 1j x” operation - a unary operation.

x.exp1j() returns a ComplexLazyTensor that encodes, symbolically, the complex exponential of 1j*x.

keops_tensordot(*args)

Tensor dot product (on KeOps internal dimensions) - a binary operation.

Parameters:

• other – a LazyTensor

• dimfa – tuple/list of int

• dimfb – tuple/list of int

• contfa – tuple/list of int listing contraction dimension of a (could be empty)

• contfb – tuple/list of int listing contraction dimension of b (could be empty)

• args – a tuple/list of int containing the graph of a permutation of the output

Returns:

a LazyTensor

__call__(*args, **kwargs)

Executes a Genred or KernelSolve call on the input data, as specified by self.formula .