Pixel operations in batches
I have a batch of depth images, shape -&gt; [B, 1, H, W]. For each pixel in each image of the batch I need to perform:
X = d * Kinverse @ [u, v, 1]   #therefore X is in R^3
where d is float tensor[0;1] representing depth at pixel u,v; Kinverse is a constant 3X3 matrix and u, v refer to the pixel column and row respectively.
Is there some way I can vectorize the operation to obtain X(u+1,v), X(u,v) and X(u,v+1) for all the images in the batch.
I eventually need to take this cross product:
{X(u+1,v) - X(u,v)} x {X(u, v+1) - X(u,v)}
Thanks for the help!

You can use torch.meshgrid to produce the u and v tensors. Once you have them, you can use torch.einsum to do the batched matrix multiplication with Kinverse.
Finally, you can use torch.cross to compute the cross product:
u, v = torch.meshgrid(*[torch.arange(s_, dtype=d.dtype, device=d.device) for s_ in d.shape[2:]])
# make a single 1x1xHxW for [u v 1] per pixel:
uv = torch.cat((u[None, None, ...], v[None, None, ...], torch.ones_like(u)[None, None, ...]), dim=1)
# compute X
X = d * torch.einsum('ij,bjhw-&gt;bihw',Kinverse,uv)
# the cross product
out = torch.cross(X[..., 1:, :-1] - X[..., :-1, :-1], X[..., :-1, 1:] - X[..., :-1, :-1], dim=1)

