How to create a 2D tensor of points with PyTorch, each dimension going from 0 to 1?

Question

I'm trying to create a 2D tensor where each dimension ranges from 0 to 1.

For a 1D tensor, I can use:

torch.arange(0, 1, 0.2)

This gives me:

tensor([0.0, 0.2, 0.4, 0.6, 0.8])

But, I want to extend this to 2D points. My desired output is [with the shape (25, 2)]:

tensor([
    [0.0, 0.0], [0.0, 0.2], [0.0, 0.4], [0.0, 0.6], [0.0, 0.8],
    [0.2, 0.0], [0.2, 0.2], [0.2, 0.4], [0.2, 0.6], [0.2, 0.8],
    [0.4, 0.0], [0.4, 0.2], [0.4, 0.4], [0.4, 0.6], [0.4, 0.8],
    [0.6, 0.0], [0.6, 0.2], [0.6, 0.4], [0.6, 0.6], [0.6, 0.8],
    [0.8, 0.0], [0.8, 0.2], [0.8, 0.4], [0.8, 0.6], [0.8, 0.8]
])

How can I achieve this using PyTorch?

Answer 1

Operation that you want is called Cartesian Product. Using torch you can achieve similar results to yours using torch.cartesian_prodd

torch.cartesian_prod(torch.arange(0, 1, 0.2), torch.arange(0, 1, 0.2))

it produces

tensor([[0.0000, 0.0000],
    [0.0000, 0.2000],
    [0.0000, 0.4000],
    [0.0000, 0.6000],
    [0.0000, 0.8000],
    [0.2000, 0.0000],
    [0.2000, 0.2000],
    [0.2000, 0.4000],
    [0.2000, 0.6000],
    [0.2000, 0.8000],
    [0.4000, 0.0000],
    [0.4000, 0.2000],
    [0.4000, 0.4000],
    [0.4000, 0.6000],
    [0.4000, 0.8000],
    [0.6000, 0.0000],
    [0.6000, 0.2000],
    [0.6000, 0.4000],
    [0.6000, 0.6000],
    [0.6000, 0.8000],
    [0.8000, 0.0000],
    [0.8000, 0.2000],
    [0.8000, 0.4000],
    [0.8000, 0.6000],
    [0.8000, 0.8000]])

Hope it helps