Basic requirement: suppose I have an array of [9, 8, 7], and I can generate two random assignment indices arrays, suppose they are [1, 1, 0] and [0, 2, 1], so the assignment results are [8, 8, 9] and [9, 7, 8], and finally I want to get the stacked result: [[8, 8, 9], [9, 7, 8]].

I can achieve the above function by the following code with a for-loop:

```
import torch
my_array = torch.Tensor([9, 8, 7])
for i in range(2): # for each layer
my_assignment_indices = torch.randint(0, 3, size=(3,)) # for each location, randomly choose a source index
current_assignment_array = my_array[my_assignment_indices].unsqueeze(0) # move (assign)
if i == 0:
result = current_assignment_array # the first layer, create result
else:
result = torch.cat([result, current_assignment_array]) # the subsequent layers, stack them together
print(result) # output my result
```

In practice, I have two parameters: N (array length), K (layer number).

Notice that each index is randomly chosen from [0, …, N-1].

Now I am confused by the following two problems:

- How can I perform the above process efficiently with very large N and K on CPU or GPU?
- How can I randomly assign K different elements to each location? The above code does not guarantee that each two indices for a location must be different.

I’m waiting for a better solution, thank you for watching my problems .