Simplify the mixing code to have only 2 dimensions instead of 3

At the moment the code always have one dimension of size 1, which is unnecessary, so we should simplify it.