Function restore

imperative/python/megengine/functional/math.py:647–650

(tensor)

    tns, ind = apply(op, reshaped)

    def restore(tensor):
        unreshaped = tensor.reshape(transposed_shape)
        untransposed = swapaxes(unreshaped, dim, -1)
        return untransposed

    return restore(tns), restore(ind)