@staticmethod
def backward(ctx, grad_output):
    from torch.distributions import OneHotCategorical
    (x, dim) = (ctx.saved_tensors, ctx.dim)
    if ctx.needs_input_grad[0]:
        return (grad_output.unsqueeze(dim).mul(OneHotCategorical(logits=x.movedim(dim, (- 1))).sample().movedim((- 1), dim)), None)
    return (None, None)
@staticmethod
def backward(ctx, grad_output):
    from torch.distributions import OneHotCategorical
    (x, dim) = (ctx.saved_tensors, ctx.dim)
    if ctx.needs_input_grad[0]:
        return (grad_output.unsqueeze(dim).mul(OneHotCategorical(logits=x.movedim(dim, (- 1))).sample().movedim((- 1), dim)), None)
    return (None, None)