jax.nn.one_hot

jax.nn.one_hot(x, num_classes, *, dtype=<class 'jax.numpy.float64'>, axis=-1)

对给定的索引进行独热编码。

输入 x 中的每个索引都被编码为一个长度为 num_classes 的零向量，其中 index 处的元素设置为 1

>>> jax.nn.one_hot(jnp.array([0, 1, 2]), 3)
Array([[1., 0., 0.],
       [0., 1., 0.],
       [0., 0., 1.]], dtype=float32)

索引超出范围 [0, num_classes) 将被编码为零

>>> jax.nn.one_hot(jnp.array([-1, 3]), 3)
Array([[0., 0., 0.],
       [0., 0., 0.]], dtype=float32)

参数：:

• x (Any) – 索引的张量。

• num_classes (int) – 独热维度中的类别数。

• dtype (Any) – 可选，返回值的浮点类型（默认 jnp.float_）。

• axis (int | AxisName) – 应该计算函数的轴或轴。

返回类型：:

数组