Wrap a static operator into a unified operator with runtime dispatching. Args: original_cpp_ext_op: Original C++ extension operator function. original_custom_op: Original custom operator function.
(original_cpp_ext_op, original_custom_op)
| 64 | |
| 65 | |
| 66 | def wrap_unified_op(original_cpp_ext_op, original_custom_op): |
| 67 | """ |
| 68 | Wrap a static operator into a unified operator with runtime dispatching. |
| 69 | Args: |
| 70 | original_cpp_ext_op: Original C++ extension operator function. |
| 71 | original_custom_op: Original custom operator function. |
| 72 | """ |
| 73 | try: |
| 74 | |
| 75 | @paddle.jit.marker.unified |
| 76 | @functools.wraps(original_custom_op) |
| 77 | def unified_op(*args, **kwargs): |
| 78 | if paddle.in_dynamic_mode(): |
| 79 | res = original_cpp_ext_op(*args, **kwargs) |
| 80 | if res is None: |
| 81 | return None |
| 82 | # TODO(DrRyanHuang): Remove this if when we align the implementation of custom op and C++ extension |
| 83 | if isinstance(res, list) and len(res) == 1: |
| 84 | return res[0] |
| 85 | return res |
| 86 | return original_custom_op(*args, **kwargs) |
| 87 | |
| 88 | except: |
| 89 | unified_op = None |
| 90 | logger.warning("Paddle version not support JIT mode.") |
| 91 | return unified_op |
| 92 | |
| 93 | |
| 94 | def preprocess_static_op(global_ns): |
no outgoing calls
no test coverage detected