(mod: tvm.IRModule)
| 49 | |
| 50 | |
| 51 | def decompose_clip(mod: tvm.IRModule) -> tvm.IRModule: |
| 52 | from tvm.relax import Expr |
| 53 | from tvm.relax.dpl import DFPattern, rewrite_call |
| 54 | from tvm.relax.dpl.pattern import is_op, wildcard |
| 55 | |
| 56 | input_pattern = wildcard() |
| 57 | min_pattern = wildcard() |
| 58 | max_pattern = wildcard() |
| 59 | pattern = is_op("relax.clip")(input_pattern, min_pattern, max_pattern) |
| 60 | |
| 61 | def _rewriter(expr: Expr, matches: dict[DFPattern, Expr]) -> Expr: # pylint: disable=unused-argument |
| 62 | dtype = matches[input_pattern].struct_info.dtype |
| 63 | return R.minimum( |
| 64 | R.maximum( |
| 65 | matches[input_pattern], |
| 66 | R.const(np.array(matches[min_pattern].value.value).astype(dtype), dtype), |
| 67 | ), |
| 68 | R.const(np.array(matches[max_pattern].value.value).astype(dtype), dtype), |
| 69 | ) |
| 70 | |
| 71 | mod["main"] = rewrite_call(pattern, _rewriter, mod["main"]) |
| 72 | return mod |
| 73 | |
| 74 | |
| 75 | def _build(mod, enable_nnapi): |
no test coverage detected
searching dependent graphs…