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


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

if get_instance: if get_instance:
return CompileOp() return CompileOp()
return CompileOp return CompileOp


Loading…
Cancel
Save