Cannot sum() over dim=0 with nbatchdim=1

[NicPy4]

I work with PyKeops through GPyTorch. I'm trying to implement a new kernel, a so called Spectral Mixture Kernel. I think that I implemented the kernel correct so far. I followed the tips in the GPyTorch Q&A. One of the last steps of calculating the covariance matrix consists of a sum reduction in dim=0. We have a LazyTensor of shape (k, n, n, d), with k representing the different mixtures, n the amount of data and d the dimensions. So the sum reduction should sum over the k-mixtures in dim=0 and result in the shape (1, n, n, d). This is the code:

from pykeops.torch.lazytensor.LazyTensor import LazyTensor
import math
import gpytorch
import pykeops
from linear_operator.operators import KernelLinearOperator
from gpytorch.kernels.spectral_mixture_kernel import SpectralMixtureKernel as GSpectralMixtureKernel
from gpytorch.kernels.keops.keops_kernel import KeOpsKernel

try:
    from pykeops.torch import LazyTensor as KEOLazyTensor

    def _covar_func(x1, x2, mixture_scales, mixture_means, mixture_weights, num_mixtures, **kwargs):
      # 'k x 1 x 1 x 1'
      mixture_weights = mixture_weights[ :, None, None, None]
      # compute distances - scaled by appropriate parameters
      x1_exp = x1 * mixture_scales
      x2_exp = x2 * mixture_scales
      x1_cos = x1 * mixture_means
      x2_cos = x2 * mixture_means

      # 'k x n x 1 x d'
      x1_exp_ = KEOLazyTensor(x1_exp[..., :, None, :])
      # 'k x 1 x n x d'
      x2_exp_ = KEOLazyTensor(x2_exp[..., None, :, :])
      # 'k x n x 1 x d'
      x1_cos_ = KEOLazyTensor(x1_cos[..., :, None, :])
      # 'k x 1 x n x d'
      x2_cos_ = KEOLazyTensor(x2_cos[..., None, :, :])
      # Compute the exponential and cosine terms
      k_exp = ((x1_exp_ - x2_exp_)**2) * (-2 * math.pi**2)
      k_cos = (x1_cos_ - x2_cos_) * (2 * math.pi)
      k = k_exp.exp() * k_cos.cos()
      # 'k x n x n x d'
      res = (k * mixture_weights).sum(dim=0)
      res = res.prod(-1)
      return res

    class SpectralMixtureKernel(KeOpsKernel, GSpectralMixtureKernel):
      '''
      Implements the Spectral Mixture Kernel using KeOps as a driver for kernel matrix multiplies.


      '''
      is_stationary = True
          

      def _nonkeops_forward(self, x1, x2, diag=False, **kwargs):
          print('we are in nonkeops')
          n, num_dims = x1.shape[-2:]
          if not num_dims == self.ard_num_dims:
              raise RuntimeError(
                  "The SpectralMixtureKernel expected the input to have {} dimensionality "
                  "(based on the ard_num_dims argument). Got {}.".format(self.ard_num_dims, num_dims)
              )

          x1_ = x1.unsqueeze(-3)
          x2_ = x2.unsqueeze(-3)

          # compute distances - scaled by appropriate parameters
          x1_exp = x1_ * self.mixture_scales
          x2_exp = x2_ * self.mixture_scales
          x1_cos = x1_ * self.mixture_means
          x2_cos = x2_ * self.mixture_means

          # Create grids
          x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp, diag=diag)
          x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos, diag=diag)

          # Compute the exponential and cosine terms
          k_exp = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi ** 2)
          k_cos = (x1_cos_ - x2_cos_).mul_(2 * math.pi)
          k = k_exp.exp_() * k_cos.cos_()

          # sum over mixtures
          mixture_weights = self.mixture_weights.view(*self.mixture_weights.shape, 1, 1)
          if not diag:
              mixture_weights = mixture_weights.unsqueeze(-2)
          res = (k * mixture_weights).sum(-3 if diag else -4)
           # Product over dimensions
          res = res.prod(-1)

          return res

      def _keops_forward(self, x1, x2, **kwargs):
          n, num_dims = x1.shape[-2:]
          if not num_dims == self.ard_num_dims:
              raise RuntimeError(
                  "The SpectralMixtureKernel expected the input to have {} dimensionality "
                  "(based on the ard_num_dims argument). Got {}.".format(self.ard_num_dims, num_dims)
              )
          x1_ = x1.unsqueeze(-3)
          x2_ = x2.unsqueeze(-3)
          # return KernelLinearOperator inst only when calculating the whole covariance matrix
          return KernelLinearOperator(x1_, x2_, covar_func=_covar_func(x1_, x2_, self.mixture_scales, self.mixture_means, self.mixture_weights, self.num_mixtures),
                                      **kwargs)
except ImportError:

    class SpectralMixtureKernel(GSpectralMixtureKernel):
        pass

But nbatchdim=1, so when I try to do the sum reduction on dim=0 in the _covar_func, I get this error:

[<ipython-input-34-71fe6c68b6cd>](https://localhost:8080/#) in _covar_func(x1, x2, mixture_scales, mixture_means, mixture_weights, num_mixtures, **kwargs)
     32       k = k_exp.exp() * k_cos.cos()
     33       # 'k x n x n x d'
---> 34       res = (k * mixture_weights).sum(dim=0)
     35       res = res.prod(-1)
     36       return res

[/usr/local/lib/python3.10/dist-packages/pykeops/common/lazy_tensor.py](https://localhost:8080/#) in sum(self, axis, dim, **kwargs)
   1879             return self.unary("Sum", dimres=1)
   1880         else:
-> 1881             return self.reduction("Sum", axis=axis, **kwargs)
   1882 
   1883     def sum_reduction(self, axis=None, dim=None, **kwargs):

[/usr/local/lib/python3.10/dist-packages/pykeops/common/lazy_tensor.py](https://localhost:8080/#) in reduction(self, reduction_op, other, opt_arg, axis, dim, call, is_complex, **kwargs)
    717             axis = dim  # NumPy uses axis, PyTorch uses dim...
    718         if axis - self.nbatchdims not in (0, 1):
--> 719             raise ValueError(
    720                 "Reductions must be called with 'axis' (or 'dim') equal to the number of batch dimensions + 0 or 1."
    721             )

ValueError: Reductions must be called with 'axis' (or 'dim') equal to the number of batch dimensions + 0 or 1.

I understand the origin of the error but I have to sum over dim=0 because I have to sum over the different mixtures and I can't find a solution so far. I tried to sum over other dimensions but then I get a wrong result.
So I hope that anyone can help me or have a tip for me how I can solve this. Appreciate it!
x1 and x2 have the shape (1, n, d).