take a 2d numpy array of category labels and turn it into a 3d one-hot numpy array

2d_to_3d.py

import numpy as np

# the 2d array of our samples,
# each component is a category label
a = np.array([[1,2,3],[4,5,6]])

# the 3d array that will be the one-hot representation
# a.max() + 1 is the number of labels we have
b = np.zeros((a.shape[0], a.shape[1], a.max() + 1))

# if you visualize this as a stack of layers,
# where each layer is a sample,
# this first index selects each layer separately
layer_idx = np.arange(a.shape[0]).reshape(a.shape[0], 1)

# this index selects each component separately
component_idx = np.tile(np.arange(a.shape[1]), (a.shape[0], 1))

# then we use `a` to select indices according to category label
b[layer_idx, component_idx, a] = 1

# voila!
print(b)

Is it possible to inverse the process ? After encoding 2D to 3D numpy array, I want to get 2 D array decoded out of 3D array encoded.

there's a shorter one

np.arange(a.max()+1) == a[...,None]

How does one reverse this?

This really helped me! Thank you