假设我有一个尺寸为x的张量[num_elem, 128]和另一个尺寸为batch的张量[num_elem, 1]。 batch中的条目是从0到511的整数，指示x中每个条目对应的批次。此外，每个批次的计数数量可能不同（例如0出现22次，而29仅出现5次）。
我想将x重塑为维度[batch_size (512), unk, 128]的张量，其中unk指batch中整数的最大数量。目前，我正在找出batch的模态频率，然后构造尺寸为x_reshaped的张量[512, mode_freq, 128]。之后，我使用for循环遍历x，将其分配给正确的批次（通过引用batch）和x_reshaped的第二维。但是，这非常慢。
还有其他更有效的方法吗？谢谢。