1# Copyright 2022 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15 16"""Register pyfunc for opaque_func_cpu_kernel""" 17 18from mindspore.ops._register_for_op import OpaquePredicateRegistry 19 20 21registered_func_name = OpaquePredicateRegistry() 22 23 24def add_opaque_predicate(fn_name, func): 25 """restore opaque predicate functions""" 26 registered_func_name.register(fn_name, func) 27 28 29def get_opaque_predicate(fn_name): 30 """get opaque predicate function by their name""" 31 return registered_func_name.get(fn_name) 32 33 34def get_func_names(): 35 """get function names""" 36 return registered_func_name.func_names 37 38 39def clean_funcs(): 40 """clean restored functions""" 41 registered_func_name.func_names = [] 42