Alright, I’ve got this fully implemented here: https://github.com/SimLeek/TorchRender/blob/master/torchrender/pixel_shader.py 23
Now I just need to get tests on it and upload to pypi.
For some reason I don’t completely understand, this setup works quite well:
self.min_bounds = [0 for _ in frame.shape]
self.max_bounds = list(frame.shape)
grid_slices = [slice(self.min_bounds[d], self.max_bounds[d]) for d in range(len(frame.shape))]
space_grid = np.mgrid[grid_slices]
space_grid.flags.writeable = False
x_tens = torch.LongTensor(space_grid[0, ...]).to(self.device)
y_tens = torch.LongTensor(space_grid[1, ...]).to(self.device)
c_tens = torch.LongTensor(space_grid[2, ...]).to(self.device)
self.x = Variable(x_tens, requires_grad=False)
self.y = Variable(y_tens, requires_grad=False)
self.c = Variable(c_tens, requires_grad=False)

then, passing x,y,c into coords, I can do something like:
>>> img = np.zeros((600, 800, 3))
>>> def fun(array, coords, finished):
...     rgb = torch.empty(array.shape).uniform_(0,1).type(torch.DoubleTensor).to(array.device)/300.0
...     trans = np.zeros_like(coords)
...     trans[0,...] = np.ones(trans.shape[1:])
...     array[coords] = (array[coords+trans] + rgb[coords])%1.0
>>> VideoHandlerThread(video_source=img, callbacks=pixel_shader(fun)).display()

Strangely, using pure numpy for the meshgrid slows things down, and using pure pytorch overloads my GPU.