| 20 | import onnx_graphsurgeon as gs |
| 21 | |
| 22 | def main(): |
| 23 | input0 = gs.Variable(name="input0", dtype=np.float32, shape=('n_rows', 8)) |
| 24 | input1 = gs.Variable(name="input1", dtype=np.float32, shape=('n_rows', 8)) |
| 25 | output = gs.Variable(name="output", dtype=np.float32, ) |
| 26 | |
| 27 | node = gs.Node(op="Concat", inputs=[input0, input1], outputs=[output], attrs={"axis": 0}) |
| 28 | |
| 29 | graph = gs.Graph(nodes=[node], inputs=[input0, input1], outputs=[output]) |
| 30 | |
| 31 | model = gs.export_onnx(graph) |
| 32 | onnx.save(model, "concat_layer.onnx") |
| 33 | |
| 34 | if __name__ == '__main__': |
| 35 | main() |