revop solver

mlrev.py

import onnx
import numpy as np

g = onnx.load("./model_dist.onnx")
nodes = g.graph.node

# remove the customop nodes
nodes.remove(nodes[14])
nodes.remove(nodes[14])
nodes.remove(nodes[14])

dims = (64, 84)

# first operation: (sin(x) + cos(x))
nodes.insert(14,onnx.helper.make_node("Sin", ["pwny_ops::customop_1_29"], ["c1_a"], "c1_a"))
nodes.insert(15, onnx.helper.make_node("Cos", ["pwny_ops::customop_1_29"], ["c1_b"], "c1_b"))
nodes.insert(16, onnx.helper.make_node("Add", ["c1_a", "c1_b"], ["customop_1_30"], "c1_c"))

# second operation: x - tanh(x)
nodes.insert(17, onnx.helper.make_node("Tanh", ["customop_1_30"], ["c2_b"], "c2_b"))
nodes.insert(18, onnx.helper.make_node("Sub", ["customop_1_30", "c2_b"], ["customop_1_31"], "c2_c"))

# third operation: x < 0 ? x / 100 : x
#   compiled as x*0.0099999998, but /100 makes more sense imo
values = np.zeros(dims).astype(np.float32)
tensor = onnx.helper.make_tensor(
        name="const_zero",
        data_type=onnx.TensorProto.FLOAT,
        dims=values.shape,
        vals=values.flatten().astype(float)
    )
nodes.insert(19, onnx.helper.make_node("Constant", [], ["c3_a"], "c3_a",value=tensor))
nodes.insert(20, onnx.helper.make_node("Less", ["customop_1_31", "c3_a"], ["c3_b"], "c3_b"))    # x<0

values = np.ones(dims).astype(np.float32) * 0.0099999998
tensor = onnx.helper.make_tensor(
        name="const_percentage",
        data_type=onnx.TensorProto.FLOAT,
        dims=values.shape,
        vals=values.flatten().astype(float)
    )
nodes.insert(21, onnx.helper.make_node("Constant", [], ["c3_c"], "c3_c",value=tensor))
nodes.insert(22, onnx.helper.make_node("Mul", ["customop_1_31", "c3_c"], ["c3_d"], "c3_d")) # x/100
nodes.insert(23, onnx.helper.make_node("Where", ["c3_b", "c3_d", "customop_1_31"], ["onnx::Gemm_32"], "c3_e")) # ternary

from base64 import b64encode
onnx.checker.check_model(g)
onnx.save(g, "newmodel.onnx")
with open("newmodel.onnx", "rb") as f:
    d = b64encode(f.read())

from pwn import *
r = remote("revop.chal.uiuc.tf", 1337)
r.sendlineafter(b"file:", d)
r.interactive()