Using numpy.where in 4D matrix to set values at the 3rd dimension?

Question

Say I have two 4D matrices

import numpy as np
from tensorflow.keras.utils import to_categorical

x = to_categorical(np.argmax(np.random.rand(2, 2, 2, 4), axis=3),