I would like to be able to somehow reorder a numpy array for efficient processing of tiles.

what I got:

>>> A = np.array([[1,2],[3,4]]).repeat(2,0).repeat(2,1)
>>> A  # image like array
array([[[1, 1, 2, 2],
        [1, 1, 2, 2]],

       [[3, 3, 4, 4],
        [3, 3, 4, 4]]])

>>> A.reshape(2,2,4)
array([[[1, 1, 2, 2],
        [1, 1, 2, 2]],

       [[3, 3, 4, 4],
        [3, 3, 4, 4]]])

what I want: X

>>> X
array([[[1, 1, 1, 1],
        [2, 2, 2, 2]],

       [[3, 3, 3, 3],
        [4, 4, 4, 4]]])

to be able to do something like:

>>> X[X.sum(2)>12] -= 1
>>> X
array([[[1, 1, 1, 1],
        [2, 2, 2, 2]],

       [[3, 3, 3, 3],
        [3, 3, 3, 3]]])

Is this possible without a slow python loop?

Bonus: Conversion back from X to A

Edit: How can I get X from A?