Best way to get the first n rows from a NX Tensor

Ops. my fault.
I have just realized that the created tensor in Nx is different to the array created in numpy. I didn’t realize it before because the tensor that I used as example was 3x3, but if I use a not squared tensor (eg. 3x2) it is evident.

x = Nx.tensor([
  [0, 1],
  [2, 3],
  [4, 5]
])

# let's get the shape
Nx.shape(x)
{3, 2}

Is not the same of:

import numpy as np
x1 = [0, 1]
x2 = [2, 3]
x3 = [4, 5]
X = np.column_stack([x1, x2, x3])

>>> X
array([[0, 2, 4],
       [1, 3, 5]])

>>> X.shape
(2, 3)

Basically, I need to transpose the tensor in Elixir to have the equivalent of the numpy array:

x = Nx.tensor([
  [0, 1],
  [2, 3],
  [4, 5]
])

# let's get the shape
Nx.shape(x)
{3, 2}

# and transpose it
Nx.transpose(x)
#Nx.Tensor<
  s64[2][3]
  [
    [0, 2, 4],
    [1, 3, 5]
  ]
>

t = Nx.tensor([
  [0, 1, 2],
  [3, 4, 5],
  [6, 7, 8]
])

t
|> Nx.slice([0, 0], [3, 2])  # slice from index (0, 0) and take 3 and 2 elements respectively along each axis
|> Nx.transpose()

returns

#Nx.Tensor<
  s64[2][3]
  [
    [0, 3, 6],
    [1, 4, 7]
  ]
>

Thank youuuu @Aetherus