import einops

# 創(chuàng)建一個(gè)形狀為(batch_size, seq_length, hidden_dim)的張量
tensor = tf.constant([[[1, 2, 3, 4],
                       [5, 6, 7, 8]],
                      [[9, 10, 11, 12],
                       [13, 14, 15, 16]]])

# 使用einops進(jìn)行維度交換和重塑
reshaped = einops.rearrange(tensor, 'batch seq dim -> batch dim seq')

print(reshaped.shape)