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)