import numpy as np
import torch


# [batch_size * num_choices, 1] == [3 * 6, 1]

torch.ones((2*3, 1))

xx = torch.arange(0, 9).view(3*3, 1)

# xx = torch.Tensor(
#     [1, 2, 3, 4, 5, 6],
#     [11, 22, 33, 44, 55, 66],
#     [111, 222, 333, 444, 555, 666]
# )

xx.view(3, 3)


xx = torch.arange(0, 12).view(3*4, 1)
xx.views()
xx[None, :].shape
xx[:, None].shape
xx.shape

# valid_len = torch.ones((12)).view(3*4, 1)
valid_len = torch.cat((torch.ones(2), torch.zeros(8))).view(3*4, 1)
valid_len = torch.Tensor([0, 1, 2, 3]).view(4, 1)
mask = torch.arange((4))[None, :]
mask.shape
valid_len.shape

mask < valid_len

xx[mask]