Note

Click here to download the full example code

# Advanced usage: Vi, Vj, Pm helpers and symbolic variables¶

This tutorial shows some advanced features of the LazyTensor class.

```
import time
import torch
from pykeops.torch import LazyTensor
use_cuda = torch.cuda.is_available()
tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
```

## The Vi, Vj, Pm decorators¶

This part presents an alternative style for using the KeOps LazyTensor wrapper, that
some users may find more convenient. The idea is to always input 2D tensors,
and use the `Vi`

, `Vj`

helpers described below to specify wether the tensor is to be
understood as indexed by i (i.e. with an equivalent shape of the form (M,1,D))
or by j (shape of the form (1,N,D)). Note that it is currently not possible to use
additional batch dimensions with this specific syntax.

Here is how it works, if we want to perform a simple gaussian convolution.

We first create the dataset using 2D tensors:

```
M, N = (100000, 200000) if use_cuda else (1000, 2000)
D = 3
x = torch.randn(M, D).type(tensor)
y = torch.randn(N, D).type(tensor)
```

Then we use `Vi`

and `Vj`

to convert to KeOps LazyTensor objects

```
from pykeops.torch import Vi, Vj
x_i = Vi(x) # (M, 1, D) LazyTensor, equivalent to LazyTensor( x[:,None,:] )
y_j = Vj(y) # (1, N, D) LazyTensor, equivalent to LazyTensor( y[None,:,:] )
```

and perform our operations:

```
D2xy = ((x_i - y_j) ** 2).sum()
gamma = D2xy.sum_reduction(dim=1)
```

Note that in the first line we used `sum`

without any axis or dim parameter.
This is equivalent to `sum(-1)`

or `sum(dim=2)`

, because
the axis parameter is set to `2`

by default. But speaking about `dim=2`

here with the `Vi`

, `Vj`

helpers could be misleading.
Similarly we used `sum_reduction`

instead of `sum`

to make it clear
that we perform a reduction, but sum and sum_reduction with `dim=0`

or `1`

are equivalent (however `sum_reduction`

with `dim=2`

is forbidden)

We have not spoken about `Pm`

yet. In fact `Pm`

is used to introduce
scalars or 1D vectors of parameters into formulas, but it is useless
in such examples because scalars, lists of scalars, 0D or 1D NumPy vectors
are automatically converted into parameters when combined with
KeOps formulas. We will have to use `Pm`

in other parts below.

## Other examples¶

All KeOps operations and reductions are available, either via operators or methods. Here are one line examples

Getting indices of closest point between x and y:

```
indmin = ((x_i - y_j) ** 2).sum().argmin(axis=0)
```

Out:

```
Compiling libKeOpstorch4bd7e3de28 in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorch4bd7e3de28:
formula: ArgMin_Reduction(Sum(Square((Var(0,3,0) - Var(1,3,1)))),1)
aliases: Var(0,3,0); Var(1,3,1);
dtype : float32
... Done.
```

Scalar product, absolute value, power operator, and a SoftMax type reduction:

```
res = (abs(x_i | y_j) ** 1.5).sumsoftmaxweight(x_i, axis=1)
```

Out:

```
Compiling libKeOpstorch332db0cd11 in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorch332db0cd11:
formula: Max_SumShiftExpWeight_Reduction(Powf(Abs((Var(0,3,0) | Var(1,3,1))), Var(2,1,2)),0,Concat(IntCst(1),Var(0,3,0)))
aliases: Var(0,3,0); Var(1,3,1); Var(2,1,2);
dtype : float32
... Done.
```

The `[]`

operator can be used to do element selection or slicing
(Elem or Extract operation in KeOps).

```
res = (x_i[:2] * y_j[2:] - x_i[2:] * y_j[:2]).sqnorm2().sum(axis=1)
```

Out:

```
Compiling libKeOpstorch5b36959de3 in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorch5b36959de3:
formula: Sum_Reduction(SqNorm2(((Extract(Var(0,3,0),0,2) * Extract(Var(1,3,1),2,1)) - (Extract(Var(0,3,0),2,1) * Extract(Var(1,3,1),0,2)))),0)
aliases: Var(0,3,0); Var(1,3,1);
dtype : float32
... Done.
```

Kernel inversion : let’s do a gaussian kernel inversion. Note that
we have to use both `Vi`

and `Vj`

helpers on the same tensor `x`

here.

```
e_i = Vi(torch.rand(M, D).type(tensor))
x_j = Vj(x)
D2xx = LazyTensor.sum((x_i - x_j) ** 2)
sigma = 0.25
Kxx = (-D2xx / sigma ** 2).exp()
res = LazyTensor.solve(Kxx, e_i, alpha=.1)
```

Out:

```
Compiling libKeOpstorchb76a69760c in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorchb76a69760c:
formula: Sum_Reduction((Exp((Minus(Sum(Square((Var(1,3,0) - Var(2,3,1))))) / Var(3,1,2))) * Var(0,3,1)),0)
aliases: Var(0,3,1); Var(1,3,0); Var(2,3,1); Var(3,1,2);
dtype : float32
... Done.
```

## Use of loops or vector operations for sums of kernels¶

Let us now perform again a kernel convolution, but replacing the gaussian kernel by a sum of 4 gaussian kernels with different sigma widths. This can be done as follows with a for loop:

```
sigmas = tensor([0.5, 1.0, 2.0, 4.0])
b_j = Vj(torch.rand(N, D).type(tensor))
Kxy = 0
for sigma in sigmas:
Kxy += LazyTensor.exp(-D2xy / sigma ** 2)
gamma = (Kxy * b_j).sum_reduction(axis=1)
```

Out:

```
Compiling libKeOpstorch828aafbac0 in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorch828aafbac0:
formula: Sum_Reduction(((((Exp((Minus(Sum(Square((Var(0,3,0) - Var(1,3,1))))) / Var(2,1,2))) + Exp((Minus(Sum(Square((Var(0,3,0) - Var(1,3,1))))) / Var(3,1,2)))) + Exp((Minus(Sum(Square((Var(0,3,0) - Var(1,3,1))))) / Var(4,1,2)))) + Exp((Minus(Sum(Square((Var(0,3,0) - Var(1,3,1))))) / Var(5,1,2)))) * Var(6,3,1)),0)
aliases: Var(0,3,0); Var(1,3,1); Var(2,1,2); Var(3,1,2); Var(4,1,2); Var(5,1,2); Var(6,3,1);
dtype : float32
... Done.
```

Note again that after the for loop, no actual computation has been performed. So we can actually build formulas with much more flexibility than with the use of Genred.

Ok, this was just to showcase the use of a for loop, however in this case there is no need for a for loop, we can do simply:

```
Kxy = LazyTensor.exp(-D2xy / sigmas ** 2).sum()
gamma = (Kxy * b_j).sum_reduction(axis=1)
```

Out:

```
Compiling libKeOpstorchaa2c0594a5 in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorchaa2c0594a5:
formula: Sum_Reduction((Sum(Exp((Minus(Sum(Square((Var(0,3,0) - Var(1,3,1))))) / Var(2,4,2)))) * Var(3,3,1)),0)
aliases: Var(0,3,0); Var(1,3,1); Var(2,4,2); Var(3,3,1);
dtype : float32
... Done.
```

This is because all operations are broadcasted, so the `/`

operation above
works and corresponds to a `./`

(scalar-vector element-wise division)

## The “no call” mode¶

When using a reduction operation, the user has the choice to actually not perform the computation directly and instead output a KeOps object which is direclty callable. This can be done using the “call=False” option

```
M, N = (100000, 200000) if use_cuda else (1000, 2000)
D, Dv = 3, 4
x = torch.randn(M, D).type(tensor)
y = torch.randn(N, D).type(tensor)
b = torch.randn(N, Dv).type(tensor)
sigmas = tensor([0.5, 1.0, 2.0, 4.0])
from pykeops.torch import Vi, Vj, Pm
x_i = Vi(x) # (M, 1, D) LazyTensor, equivalent to LazyTensor( x[:,None,:] )
y_j = Vj(y) # (1, N, D) LazyTensor, equivalent to LazyTensor( y[None,:,:] )
b_j = Vj(b) # (1, N, D) LazyTensor, equivalent to LazyTensor( b[None,:,:] )
D2xy = ((x_i - y_j) ** 2).sum(-1)
Kxy = LazyTensor.exp(-D2xy / sigmas ** 2).sum()
gammafun = (Kxy * b_j).sum_reduction(axis=1, call=False)
```

Here gammafun is a function and can be evaluated later

```
gamma = gammafun()
```

Out:

```
Compiling libKeOpstorch2119901533 in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorch2119901533:
formula: Sum_Reduction((Sum(Exp((Minus(Sum(Square((Var(0,3,0) - Var(1,3,1))))) / Var(2,4,2)))) * Var(3,4,1)),0)
aliases: Var(0,3,0); Var(1,3,1); Var(2,4,2); Var(3,4,1);
dtype : float32
... Done.
```

This is usefull in order to avoid the small overhead caused by using the container syntax inside loops if one wants to perform a large number of times the same reduction. Here is an example where we compare the two approaches on small size data to see the effect of the overhead

```
M, N = 50, 30
x = torch.rand(M, D).type(tensor)
y = torch.rand(N, D).type(tensor)
beta = torch.rand(N, Dv).type(tensor)
xi, yj, bj = Vi(x), Vj(y), Vj(beta)
dxy2 = LazyTensor.sum((xi - yj) ** 2)
Niter = 1000
start = time.time()
for k in range(Niter):
Kxyb = LazyTensor.exp(-dxy2 / sigmas ** 2).sum() * bj
gamma = Kxyb.sum_reduction(axis=1)
end = time.time()
print('Timing for {} iterations: {:.5f}s = {} x {:.5f}s'.format(
Niter, end - start, Niter, (end - start) / Niter))
start = time.time()
Kxyb = LazyTensor.exp(-dxy2 / sigmas ** 2).sum() * bj
gammafun = Kxyb.sum_reduction(axis=1, call=False)
for k in range(Niter):
gamma = gammafun()
end = time.time()
print('Timing for {} iterations: {:.5f}s = {} x {:.5f}s'.format(
Niter, end - start, Niter, (end - start) / Niter))
```

Out:

```
Timing for 1000 iterations: 0.50551s = 1000 x 0.00051s
Timing for 1000 iterations: 0.35033s = 1000 x 0.00035s
```

Of course this means the user has to perform in-place operations
over tensors `x`

, `y,`

`beta`

inside the loop, otherwise the result of the
call to `gammafun`

will always be the same. This is not very convenient,
so we provide also a “symbolic variables” syntax.

## Using “symbolic” variables in formulas¶

Instead of inputing tensors to the `Vi`

, `Vj`

, `Pm`

helpers, one may specify
the variables as symbolic, providing an index and a dimension:

```
xi = Vi(0, D)
yj = Vj(1, D)
bj = Vj(2, Dv)
Sigmas = Pm(3, 4)
```

Now we build the formula as before

```
dxy2 = LazyTensor.sum((xi - yj) ** 2)
Kxyb = LazyTensor.exp(-dxy2 / Sigmas ** 2).sum() * bj
gammafun = Kxyb.sum_reduction(axis=1)
```

Note that we did not have to specify `call=False`

because since the
variables are symbolic, no computation can be done of course. So the
ouput is automatically a function. We can evaluate it by providing the
arguments in the order specified by the index argument given to `Vi`

, `Vj`

, `Pm`

:

```
gamma = gammafun(x, y, beta, sigmas)
```

Out:

```
Compiling libKeOpstorchbc0e6568c2 in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorchbc0e6568c2:
formula: Sum_Reduction((Sum(Exp((Minus(Sum(Square((Var(0,3,0) - Var(1,3,1))))) / Square(Var(3,4,2))))) * Var(2,4,1)),0)
aliases: Var(0,3,0); Var(1,3,1); Var(2,4,1); Var(3,4,2);
dtype : float32
... Done.
```

Symbolic and non symbolic variables can be mixed. For example if we want
to fix `x`

, `beta`

and `sigmas`

in the previous example and make the reduction
a function of `y`

only we can write:

```
xi = Vi(x)
yj = Vj(0, D)
bj = Vj(beta)
dxy2 = LazyTensor.sum((xi - yj) ** 2)
Kxyb = LazyTensor.exp(-dxy2 / sigmas ** 2).sum() * bj
gammafun = Kxyb.sum_reduction(axis=1)
print(gammafun)
gamma = gammafun(y)
```

Out:

```
KeOps LazyTensor
formula: (Sum(Exp((Minus(Sum(Square((Var(1,3,0) - Var(0,3,1))))) / Var(2,4,2)))) * Var(3,4,1))
symbolic variables: Var(0, 3, 1)
shape: (50, 30, 4)
reduction: Sum (axis=1)
Compiling libKeOpstorch33b4ddd7d3 in /home/bcharlier/tmp/libkeops/pykeops/common/../build//build-libKeOpstorch33b4ddd7d3:
formula: Sum_Reduction((Sum(Exp((Minus(Sum(Square((Var(1,3,0) - Var(0,3,1))))) / Var(2,4,2)))) * Var(3,4,1)),0)
aliases: Var(0,3,1); Var(1,3,0); Var(2,4,2); Var(3,4,1);
dtype : float32
... Done.
```

**Total running time of the script:** ( 3 minutes 1.533 seconds)