This can probably wait until the implementation is ported to KA in GPUArrays but I'm opening this to keep track somewhere.