Masking only works for convcnp

[DrJonnyT]

Based on the masking example, this works fine:

import lab as B
import torch

import neuralprocesses.torch as nps

cnp = nps.construct_convgnp(
    dim_x=2,
    dim_yc=(1, 1),  # Two context sets, both with one channel
    dim_yt=1, 
)

# Construct two sample context sets with one on a grid.
xc = B.randn(torch.float32, 1, 2, 20)
yc = B.randn(torch.float32, 1, 1, 20)
xc_grid = (B.randn(torch.float32, 1, 1, 10), B.randn(torch.float32, 1, 1, 15))
yc_grid = B.randn(torch.float32, 1, 1, 10, 15)

# Contruct sample target inputs
xt = B.randn(torch.float32, 1, 2, 50)


mask = B.ones(torch.float32, 1, 1, *B.shape(yc_grid, 2, 3))

mask[:, :, 5, 5] = 0

pred = cnp([(xc, yc), (xc_grid, nps.Masked(yc_grid, mask))], xt)

but if I replace construct_convgnp with construct_gnp or construct_agnp I get an error on the last line:
RuntimeError: Dispatched to fallback implementation for code, but specialised implementation are available.

Hardware: Mac M2, happens on both CPU and MPS GPU