torch.meshgrid
Take N tensors, each of which can be either scalar or 1-dimensional vector, and create N N-dimensional grids, where the i-th grid is defined by expanding the i-th input over dimensions defined by other inputs.
x = torch.tensor([1, 2, 3])
y = torch.tensor([4, 5, 6])
grid_x, grid_y = torch.meshgrid(x, y)
grid_x
tensor([[1, 1, 1],
[2, 2, 2],
[3, 3, 3]])
grid_y
tensor([[4, 5, 6],
[4, 5, 6],
[4, 5, 6]])