import onnx
import onnx_ir as ir

from onnxscript.rewriter import pattern, rewrite


def custom_rewrite(model):
    def target_pattern(op, x):
        return op.Relu(x, _outputs=["relu"])

    def replacement_pattern(op, x, relu):
        ir_node = relu.producer()
        return op.Clip(x, name=ir_node.name)
       """Workaround
        out = op.Clip(x)
        clip_node = out.producer()
        clip_node.name = ir_node.name
        """

    rules = [pattern.RewriteRule(target_pattern, replacement_pattern)]
    model = rewrite(model, rules)
    return model


ir_model = ir.from_onnx_text("""
            < ir_version: 10, opset_import: ["" : 20] >
            test_model (float[N, 16, 16] X) => (float [N, ?, ?] Y)
            {
                Y = Relu(X)
            }
        """)
ir_model.graph[0].name = "relu_1"
ir_model = custom_rewrite(ir_model)
proto = ir.to_proto(ir_model)
onnx.save(proto, "tmp.onnx")