One-hot representation of a matrix in numpy

Question

What is the easiest/smartest way of going from a matrix of values to one hot representation of the same thing in 3d tensor? For example if the matrix is the index after argmax in a tensor like:

indices=numpy.argmax(mytensor,axis=2)

Where tensor is 3D [x,y,z] and indices will naturally be [x,y]. Now you want to go to a 3D [x,y,z] tensor that has 1s in the place of maxes in axis=2 and 0 in any other place.

P.S. I know the answer for vector to 1-hot matrix, but this is matrix to 1-hot tensor.

Answer 1

One of the perfect setups to use broadcasting -

indices[...,None] == np.arange(mytensor.shape[-1])

If you need in ints of 0s and 1s, append with .astype(int)