diff --git a/imperative/python/megengine/functional/math.py b/imperative/python/megengine/functional/math.py index 0053a3825e1ad270f81ad36be821f56cc8568013..f93e5cee968f077892a82280e881b10cd8aec498 100644 --- a/imperative/python/megengine/functional/math.py +++ b/imperative/python/megengine/functional/math.py @@ -29,6 +29,7 @@ __all__ = [ "dot", "isinf", "isnan", + "matinv", "matmul", "max", "mean", @@ -729,6 +730,38 @@ def topk( return tns, ind +def matinv(inp: Tensor) -> Tensor: + """ + Computes the inverse of a batch of matrices; input must has shape [..., n, n]. + + :param inp: input tensor. + :return: output tensor. + + Examples: + + .. testcode:: + + import numpy as np + from megengine import tensor + import megengine.functional as F + + data = tensor([[1.0, 0.0], [1.0, 1.0]]) + out = F.matinv(data) + print(out.numpy()) + + Outputs: + + .. testoutput:: + + [[ 1. 0.] + [-1. 1.]] + + """ + + (result,) = apply(builtin.MatrixInverse(), inp) + return result + + def matmul( inp1: Tensor, inp2: Tensor, diff --git a/imperative/python/megengine/functional/nn.py b/imperative/python/megengine/functional/nn.py index fb6d4ae5d1b1e57fa3946ff0424ffabadc1163ec..91f484cefaf748941f64f9e9cc061ebfdfc5a202 100644 --- a/imperative/python/megengine/functional/nn.py +++ b/imperative/python/megengine/functional/nn.py @@ -53,7 +53,6 @@ __all__ = [ "logsigmoid", "logsumexp", "logsoftmax", - "matinv", "max_pool2d", "one_hot", "prelu", @@ -1183,38 +1182,6 @@ def remap( return result -def matinv(inp: Tensor) -> Tensor: - """ - Computes the inverse of a batch of matrices; input must has shape [..., n, n]. - - :param inp: input tensor. - :return: output tensor. - - Examples: - - .. testcode:: - - import numpy as np - from megengine import tensor - import megengine.functional as F - - data = tensor([[1.0, 0.0], [1.0, 1.0]]) - out = F.matinv(data) - print(out.numpy()) - - Outputs: - - .. testoutput:: - - [[ 1. 0.] - [-1. 1.]] - - """ - - (result,) = apply(builtin.MatrixInverse(), inp) - return result - - def interpolate( inp: Tensor, size: Optional[Union[int, Tuple[int, int]]] = None,