jax.nn.one_hot

jax.nn.one_hot(x, num_classes, *, dtype=None, axis=-1)

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

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

>>> 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 | None) – 可选，返回值的浮点数数据类型（默认为 jnp.float_）。

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

返回类型:

Array