@staticmethod
def backward(ctx, grad_output):
    (input, axis, ret) = ctx.saved_tensors
    shape = [(1 if ((axis is None) or (i in axis)) else input.shape[i]) for i in range(len(input.shape))]
    ret2 = (input == ret.reshape(shape))
    div = ret2.sum(axis=(None if (axis is None) else tuple(axis)), keepdims=True)
    return ((ret2 * grad_output.reshape(shape)) / div)
@staticmethod
def backward(ctx, grad_output):
    (input, axis, ret) = ctx.saved_tensors
    shape = [(1 if ((axis is None) or (i in axis)) else input.shape[i]) for i in range(len(input.shape))]
    ret2 = (input == ret.reshape(shape))
    div = ret2.sum(axis=(None if (axis is None) else tuple(axis)), keepdims=True)
    return ((ret2 * grad_output.reshape(shape)) / div)