sorting along multidimensional index in numpy

I have a three dimensional array, that I need to sort. I also have a 2 dimensional index that I got with a lexsort that I would like to use for sorting. I cannot find how to apply the index to the array without extra dimensions sprouting up...

import numpy as np

a = np.array([
    [
        [1, 1, 1],
        [0, 1, 2],
        [3, 1, 1],
        [-1, 1, 1],
        [0, 2, 2],
        [1, 2, 2],
    ],[
        [-1, 1, 1],
        [0, 1, 2],
        [-3, 1, 1],
        [1, 1, 1],
        [0, -2, 2],
        [-1, -2, 2],
    ],[
        [-1, 1, 1],
        [-0, 1, -2],
        [-3, 1, 1],
        [1, 1, 1],
        [-0, -2, -2],
        [-1, -2, -2],
    ],
])
sorted_index = np.lexsort((a[:,:,0], a[:,:,1],a[:,:,2]),axis=1)

sorted_index is

[[3 0 2 1 4 5], [2 0 3 5 4 1], [5 4 1 2 0 3]]

the end result should be

a_sorted = np.array([
    [
        [-1, 1, 1],
        [1, 1, 1],
        [3, 1, 1],
        [0, 1, 2],
        [0, 2, 2],
        [1, 2, 2],
    ],[
        [-3, 1, 1],
        [-1, 1, 1],
        [1, 1, 1],
        [-1, -2, 2],
        [0, -2, 2],
        [0, 1, 2],
    ],[
        [-1, -2, -2],
        [0, -2, -2],
        [0, 1, -2],  
        [-3, 1, 1],
        [-1, 1, 1],
        [1, 1, 1],
    ],
])

I tried.

a[sorted_index]
a[sorted_index, :]
a[:, sorted_index]

I am about to loop through the thing, but I would like to do it properly.