Browse Source

PyNative peformance optimize. Avoid calls constexpr InferShape

pull/14797/head
caifubi 4 years ago
parent
commit
dec192ab5e
1 changed files with 6 additions and 0 deletions
  1. +6
    -0
      mindspore/ops/primitive.py

+ 6
- 0
mindspore/ops/primitive.py View File

@@ -537,6 +537,9 @@ def constexpr(fn=None, get_instance=True, name=None):

def deco(fn):
class CompileOp(PrimitiveWithInfer):
"""
CompileOp is a temporary operator used to execute the constexpr function.
"""
def __init__(self):
op_name = name if name else fn.__name__
PrimitiveWithInfer.__init__(self, op_name)
@@ -545,6 +548,9 @@ def constexpr(fn=None, get_instance=True, name=None):
def infer_value(self, *args):
return fn(*args)

def __call__(self, *args, **kwargs):
return fn(*args)

if get_instance:
return CompileOp()
return CompileOp


Loading…
Cancel
Save