falkon.mmv_ops

The algorithms to compute kernels and kernel-vector products blockwise on GPUs and CPU. The algorithms in this module are kernel agnostic. Refer to falkon.kernels for the actual kernel implementations.

The KeOps wrapper only supports the mmv operation (kernel-vector products). The matrix-multiplication implementations instead support three different operations:

  • mm which calculates the full kernel

  • mmv which calculates kernel-vector products

  • dmmv which calculates double kernel-vector products (which are operations like \(K^\top (K v)\) where \(K\) is a kernel matrix and \(v\) is a vector).

run_keops_mmv

A thin wrapper to KeOps is provided to allow for block-splitting and multiple GPU usage. This only supports kernel-vector products.

falkon.mmv_ops.keops.run_keops_mmv(X1: torch.Tensor, X2: torch.Tensor, v: torch.Tensor, other_vars: List[torch.Tensor], out: Optional[torch.Tensor], formula: str, aliases: List[str], axis: int, reduction: str = 'Sum', opt: Optional[falkon.options.FalkonOptions] = None)torch.Tensor

fmm_cpu

Blocked kernel calculation on the CPU.

falkon.mmv_ops.fmm_cpu.fmm_cpu(X1: torch.Tensor, X2: torch.Tensor, kernel: falkon.kernels.kernel.Kernel, out: Optional[torch.Tensor], opt: falkon.options.BaseOptions)torch.Tensor

Compute kernel value on matrices X1 and X2: out = kernel(X1, X2)

Parameters
  • X1 – [N, D] array

  • X2 – [M, D] array

  • kernel – Class representing the desired kernel function

  • out – Array for storing the kernel output. If None, will be allocated within the function.

  • opt – Basic options dictionary, used for determining available memory. Additionally, the no_single_kernel option is used to determine the accumulator data type.

Returns

out – [N, M] array. The kernel between X1 and X2.

fmm_cpu_sparse

Blocked kernel calculation on the CPU for sparse datasets.

falkon.mmv_ops.fmm_cpu.fmm_cpu_sparse(X1: falkon.sparse.sparse_tensor.SparseTensor, X2: falkon.sparse.sparse_tensor.SparseTensor, kernel: falkon.kernels.kernel.Kernel, out: Optional[torch.Tensor], opt: falkon.options.BaseOptions)torch.Tensor

fmm_cuda

Blocked kernel calculation on GPUs.

falkon.mmv_ops.fmm_cuda.fmm_cuda(X1: torch.Tensor, X2: torch.Tensor, kernel: falkon.kernels.kernel.Kernel, out: Optional[torch.Tensor] = None, opt: Optional[falkon.options.BaseOptions] = None)torch.Tensor

performs fnc(X1*X2’, X1, X2) in blocks on multiple GPUs

fmm_cuda_sparse

Blocked kernel calculation on GPUs for sparse datasets.

falkon.mmv_ops.fmm_cuda.fmm_cuda_sparse(X1: falkon.sparse.sparse_tensor.SparseTensor, X2: falkon.sparse.sparse_tensor.SparseTensor, kernel: falkon.kernels.kernel.Kernel, out: Optional[torch.Tensor] = None, opt: Optional[falkon.options.BaseOptions] = None)torch.Tensor

fmmv_cpu

Blocked kernel-vector products on CPU.

falkon.mmv_ops.fmmv_cpu.fmmv_cpu(X1, X2, v, kernel, out, opt)

Blockwise kernel-vector product

This function computes kernel(X1, X2) @ v in a blockwise fashion, to avoid having the whole N*M kernel matrix in memory at once. Note that while the principle is that of matrix-vector product, v can have more than one column.

Parameters
  • X1 – [N, D] array

  • X2 – [M, D] array

  • v – [M, T] array

  • kernel – Class representing the desired kernel function

  • out (torch.Tensor or None) – [N, T] array for storing the kernel-vector product output. If None, will be allocated within the function.

  • opt – Basic options dictionary, used for determining available memory.

fmmv_cpu_sparse

Blocked kernel-vector products on CPU for sparse datasets.

falkon.mmv_ops.fmmv_cpu.fmmv_cpu_sparse(X1: falkon.sparse.sparse_tensor.SparseTensor, X2: falkon.sparse.sparse_tensor.SparseTensor, v: torch.Tensor, kernel: falkon.kernels.kernel.Kernel, out: Optional[torch.Tensor], opt: falkon.options.BaseOptions)

fdmmv_cpu

Blocked double kernel-vector products on CPU.

falkon.mmv_ops.fmmv_cpu.fdmmv_cpu(X1, X2, v, w, kernel, out, opt)

Calculate a double kernel-vector product.

This function computes the following quantity: kernel(X1, X2).T @ (kernel(X1, X2) @ v + w) Where one of v or w can be empty. All arrays passed to this function must be 2-dimensional, although the second dimension can be unitary.

The expression is not computed directly. We separate the computation into smaller blocks so as to reduce the total memory consumption (the large N*M kernel matrix is never wholly stored in RAM.)

Parameters
  • X1 – [N, D] array

  • X2 – [M, D] array

  • v (torch.Tensor or None) – [M, T] array. But note that at least one of v or w must be specified.

  • w (torch.Tensor or None) – [N, T] array. But note that at least one of v or w must be specified.

  • kernel – Class representing the desired kernel function

  • out (torch.Tensor or None) – [M, T] array for storing the kernel-vector product output. If None, will be allocated within the function.

  • opt – Basic options dictionary, used for determining available memory.

fdmmv_cpu_sparse

Blocked double kernel-vector products on CPU for sparse datasets.

falkon.mmv_ops.fmmv_cpu.fdmmv_cpu_sparse(X1: falkon.sparse.sparse_tensor.SparseTensor, X2: falkon.sparse.sparse_tensor.SparseTensor, v: Optional[torch.Tensor], w: Optional[torch.Tensor], kernel, out: Optional[torch.Tensor] = None, opt: Optional[falkon.options.BaseOptions] = None)

fmmv_cuda

Blocked kernel-vector products on GPUs.

falkon.mmv_ops.fmmv_cuda.fmmv_cuda(X1: torch.Tensor, X2: torch.Tensor, v: torch.Tensor, kernel, out: Optional[torch.Tensor] = None, opt: Optional[falkon.options.BaseOptions] = None)torch.Tensor

X1 : N x D X2 : M x D v : M x T

performs fnc(X1*X2’, X1, X2) * v : N x T in blocks on multiple GPUs

fmmv_cuda_sparse

Blocked kernel-vector products on GPUs for sparse datasets.

falkon.mmv_ops.fmmv_cuda.fmmv_cuda_sparse(X1: falkon.sparse.sparse_tensor.SparseTensor, X2: falkon.sparse.sparse_tensor.SparseTensor, v: torch.Tensor, kernel, out: Optional[torch.Tensor] = None, opt: Optional[falkon.options.BaseOptions] = None)torch.Tensor

fdmmv_cuda

Blocked double kernel-vector products on GPUs.

falkon.mmv_ops.fmmv_cuda.fdmmv_cuda(X1: torch.Tensor, X2: torch.Tensor, v: Optional[torch.Tensor], w: Optional[torch.Tensor], kernel, out: Optional[torch.Tensor] = None, opt: Optional[falkon.options.BaseOptions] = None)torch.Tensor

X1 : N x D X2 : M x D v : M x T w : N x T

performs fnc(X1*X2’, X1, X2)’ * ( fnc(X1*X2’, X1, X2) * v + w ) : M x T in blocks on multiple GPUs

Assume all inputs have the same data type

fdmmv_cuda_sparse

Blocked double kernel-vector products on GPUs for sparse datasets.

falkon.mmv_ops.fmmv_cuda.fdmmv_cuda_sparse(X1: falkon.sparse.sparse_tensor.SparseTensor, X2: falkon.sparse.sparse_tensor.SparseTensor, v: Optional[torch.Tensor], w: Optional[torch.Tensor], kernel, out: Optional[torch.Tensor] = None, opt: Optional[falkon.options.BaseOptions] = None)torch.Tensor

incore_fmmv

falkon.mmv_ops.fmmv_incore.incore_fmmv(mat: torch.Tensor, vec: torch.Tensor, out: Optional[torch.Tensor] = None, transpose: bool = False, opt: Optional[falkon.options.FalkonOptions] = None)torch.Tensor

incore_fdmmv

falkon.mmv_ops.fmmv_incore.incore_fdmmv(mat: torch.Tensor, vec: torch.Tensor, out: Optional[torch.Tensor] = None, opt: Optional[falkon.options.FalkonOptions] = None)torch.Tensor