mindspore.ops.gumbel_softmax

mindspore.ops.gumbel_softmax(logits, tau=1, hard=False, dim=- 1)

返回Gumbel-Softmax分布的Tensor，在 hard = True 的时候，返回one-hot形式的离散型Tensor，hard = False 时返回在dim维进行过softmax的Tensor。

参数：

• logits (Tensor) - 输入，是一个非标准化的对数概率分布。 只支持float16和float32。

• tau (float) - 标量温度, 正数。默认值：1.0。

• hard (bool) - 为True时返回one-hot离散型Tensor，可反向求导。默认值：False。

• dim (int) - 给softmax使用的参数，在dim维上做softmax操作。默认值：-1。

返回：

Tensor，shape与dtype和输入 logits 相同。

异常：

• TypeError - logits 不是Tensor。

• TypeError - logits 不是float16或float32。

• TypeError - tau 不是float。

• TypeError - hard 不是bool。

• TypeError - dim 不是int。

• ValueError - tau 不是正数。

支持平台：

Ascend GPU CPU

样例：

>>> input_x = Tensor(np.array([[-0.1, 0.3, 3.6], [0.4, 0.5, -3.2]]), mindspore.float32)
>>> output = ops.gumbel_softmax(input_x, 1.0, True, -1)
>>> print(output.shape)
(2, 3)