onnxscript_linear_function_matmul_add_model_export.py

1import onnx
2from onnxscript import script
3from onnxscript.onnx_opset import ai_onnx as op
4
5# Define an ONNX model using the @script decorator
6@script()
7def linear(A: float, X: float, B: float) -> float:
8    return op.MatMul(A, X) + B
9
10# Convert the script to an ONNX ModelProto
11onnx_model = linear.to_model_proto()
12
13# Save the model to a file
14onnx.save(onnx_model, "linear.onnx")
15
16# Use the model as a function (eager execution)
17import numpy as np
18a = np.array([[1, 2], [3, 4]], dtype=np.float32)
19x = np.array([[5, 6], [7, 8]], dtype=np.float32)
20b = np.array([[1, 1], [1, 1]], dtype=np.float32)
21
22result = linear(a, x, b)
23print(result)