TensorFlow, tf.one_hot why the shape of output is defined by the value of axis?

Question

I read the docs of tf.one_hot and found that

... . The new axis is created at dimension axis (default: the new axis is appended at the end).

What is The new axis?

If indices is a vector of length features, the output shape will be:

features x depth if axis == -1

depth x features if axis == 0

If indices is a matrix (batch) with shape [batch, features], the output shape will be:

batch x features x depth if axis == -1

batch x depth x features if axis == 1

depth x batch x features if axis == 0

Why the shape of output is defined by axis?

Answer 1

tf.one_hot() transforms a list of indices (e.g. [0, 2, 1]) and transforms it into a list of one-hot vectors of length depth.

For instance, if depth = 3,

• index 0 in the input will be replaced by [1, 0, 0]
• index 1 in the input will be replaced by [0, 1, 0]
• index 2 in the input will be replaced by [0, 0, 1]

So [0, 2, 1] would be encoded as [[1, 0, 0], [0, 0, 1], [0, 1, 0]]

As you can see, the output has one more dimension than the input (since each index is replaced by a vector).

By default (and what you usually need), the new dimension is created as the last one, so if your input is of shape (d1, d2, .., dn), your output will be of shape (d1, d2, .., dn, depth). But if you change the input parameter axis, you may choose to put the new dimension elsewhere, for instance if axis=0 your output will be of shape (depth, d1, d2, .., dn).

Changing the order of the dimensions is basically the n-dimensional version of transposing: you have the same data, but switch the order of the indices to access them (equivalent to switching the columns and the rows in a 2D matrix).