RBFKernelGrad supports multibatch
@@ -51,6 +51,7 @@ class RBFKernelGrad(RBFKernel):
"""
def forward(self, x1, x2, diag=False, **params):
batch_shape = x1.shape[:-2]
+ n_batch_dims = len(batch_shape)
n1, d = x1.shape[-2:]
n2 = x2.shape[-2]
@@ -63,7 +64,7 @@ class RBFKernelGrad(RBFKernel):
# Form all possible rank-1 products for the gradient and Hessian blocks
outer = x1_.view(*batch_shape, n1, 1, d) - x2_.view(*batch_shape, 1, n2, d)
- outer = outer / self.lengthscale
+ outer = outer / self.lengthscale.unsqueeze(-2)
outer = torch.transpose(outer, -1, -2).contiguous()
# 1) Kernel block
@@ -79,21 +80,21 @@ class RBFKernelGrad(RBFKernel):
# 2) First gradient block
outer1 = outer.view(*batch_shape, n1, n2 * d)
- K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])
+ K[..., :n1, n2:] = outer1 * K_11.repeat([*([1] * (n_batch_dims + 1)), d])
# 3) Second gradient block
outer2 = outer.transpose(-1, -3).contiguous().view(*batch_shape, n2, n1 * d)
outer2 = outer2.transpose(-1, -2)
- K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])
+ K[..., n1:, :n2] = -outer2 * K_11.repeat([*([1] * n_batch_dims), d, 1])
# 4) Hessian block
- outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])
+ outer3 = outer1.repeat([*([1] * n_batch_dims), d, 1]) * outer2.repeat([*([1] * (n_batch_dims + 1)), d])
kp = KroneckerProductLazyTensor(
- torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow_(2),
+ torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow(2),
torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1)
)
chain_rule = kp.evaluate() - outer3
- K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])
+ K[..., n1:, n2:] = chain_rule * K_11.repeat([*([1] * n_batch_dims), d, d])
# Symmetrize for stability
if n1 == n2 and torch.eq(x1, x2).all():
@@ -123,7 +124,4 @@ class RBFKernelGrad(RBFKernel):
`d` dimensions, RBFKernelGrad returns an `n_1(d+1) x n_2(d+1)` kernel matrix.
"""
non_batch_size = ((x1.size(-1) + 1) * x1.size(-2), (x2.size(-1) + 1) * x2.size(-2))
- if x1.ndimension() == 3:
- return torch.Size((x1.size(0),) + non_batch_size)
- else:
- return torch.Size(non_batch_size)
+ return torch.Size((*x1.shape[:-2], *non_batch_size))
