You need to keep the 2s next to the original dims, so the latter variant. Instead of permuting, I’d probably either sample the two independently (if the desired distribution is just the product) or do the //2 and %2 calclations to get the indices from values 0…3. (Numpy has unravel, but that is CPU only and a bit overkill here.)
Best regards
Thomas