mindspore.ms_function
- mindspore.ms_function(fn=None, obj=None, input_signature=None)[source]
Create a callable MindSpore graph from a Python function.
This allows the MindSpore runtime to apply optimizations based on graph.
- Parameters
fn (Function) – The Python function that will be run as a graph. Default: None.
obj (Object) – The Python object is used to distinguish the compiled function. Default: None.
input_signature (Tensor) – The Tensor which describes the input arguments. The shape and dtype of the Tensor will be supplied to this function. If input_signature is specified, each input to fn must be a Tensor. And the input parameters of fn cannot accept **kwargs. The shape and dtype of actual inputs should keep the same as input_signature. Otherwise, TypeError will be raised. Default: None.
- Returns
Function, if fn is not None, returns a callable function that will execute the compiled function; If fn is None, returns a decorator and when this decorator invokes with a single fn argument, the callable function is equal to the case when fn is not None.
- Supported Platforms:
Ascend
GPU
CPU
Examples
>>> import numpy as np >>> from mindspore import Tensor >>> from mindspore import ms_function ... >>> x = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32)) >>> y = Tensor(np.ones([1, 1, 3, 3]).astype(np.float32)) ... >>> # create a callable MindSpore graph by calling ms_function >>> def tensor_add(x, y): ... z = x + y ... return z ... >>> tensor_add_graph = ms_function(fn=tensor_add) >>> out = tensor_add_graph(x, y) ... >>> # create a callable MindSpore graph through decorator @ms_function >>> @ms_function ... def tensor_add_with_dec(x, y): ... z = x + y ... return z ... >>> out = tensor_add_with_dec(x, y) ... >>> # create a callable MindSpore graph through decorator @ms_function with input_signature parameter >>> @ms_function(input_signature=(Tensor(np.ones([1, 1, 3, 3]).astype(np.float32)), ... Tensor(np.ones([1, 1, 3, 3]).astype(np.float32)))) ... def tensor_add_with_sig(x, y): ... z = x + y ... return z ... >>> out = tensor_add_with_sig(x, y)