importtorchimporttorch.nnclassMyModel(torch.nn.Module):def__init__(self)->None:super().__init__()self.linear=torch.nn.Linear(2,2)defforward(self,x,bias=None):out=self.linear(x)out=out+biasreturnoutmodel=MyModel()kwargs={"bias":3.}inputs=(torch.randn(2,2,2),)export_output=torch.onnx.dynamo_export(model,inputs,**kwargs)export_output.save("my_simple_model.onnx")

fromtypingimportDict,Tupleimporttorchimporttorch.onnxdeffunc_with_nested_input_structure(x_dict:Dict[str,torch.Tensor],y_tuple:Tuple[torch.Tensor,Tuple[torch.Tensor,torch.Tensor]],):if"a"inx_dict:x=x_dict["a"]elif"b"inx_dict:x=x_dict["b"]else:x=torch.randn(3)y1,(y2,y3)=y_tuplereturnx+y1+y2+y3x_dict={"a":torch.tensor(1.)}y_tuple=(torch.tensor(2.),(torch.tensor(3.),torch.tensor(4.)))export_output=torch.onnx.dynamo_export(func_with_nested_input_structure,x_dict,y_tuple)print(export_output.adapt_torch_inputs_to_onnx(x_dict,y_tuple))

importonnx# We use ONNX opset 15 to define the function below.fromonnxscriptimportFLOATfromonnxscriptimportopset15asopfromonnxscriptimportscript# We use the script decorator to indicate that# this is meant to be translated to ONNX.@script()defonnx_hardmax(X,axis:int):"""Hardmax is similar to ArgMax, with the result being encoded OneHot style."""# The type annotation on X indicates that it is a float tensor of# unknown rank. The type annotation on axis indicates that it will# be treated as an int attribute in ONNX.## Invoke ONNX opset 15 op ArgMax.# Use unnamed arguments for ONNX input parameters, and named# arguments for ONNX attribute parameters.argmax=op.ArgMax(X,axis=axis,keepdims=False)xshape=op.Shape(X,start=axis)# use the Constant operator to create constant tensorszero=op.Constant(value_ints=[0])depth=op.GatherElements(xshape,zero)empty_shape=op.Constant(value_ints=[0])depth=op.Reshape(depth,empty_shape)values=op.Constant(value_ints=[0,1])cast_values=op.CastLike(values,X)returnop.OneHot(argmax,depth,cast_values,axis=axis)# We use the script decorator to indicate that# this is meant to be translated to ONNX.@script()defsample_model(X:FLOAT[64,128],Wt:FLOAT[128,10],Bias:FLOAT[10])->FLOAT[64,10]:matmul=op.MatMul(X,Wt)+Biasreturnonnx_hardmax(matmul,axis=1)# onnx_model is an in-memory ModelProtoonnx_model=sample_model.to_model_proto()# Save the ONNX model at a given pathonnx.save(onnx_model,"sample_model.onnx")# Check the modeltry:onnx.checker.check_model(onnx_model)exceptonnx.checker.ValidationErrorase:print(f"The model is invalid:{e}")else:print("The model is valid!")

importnumpyasnpv=np.array([[0,1],[2,3]],dtype=np.float32)result=Hardmax(v)

importonnxfromspoximportargument,build,Tensor,Var# Import operators from the ai.onnx domain at version 17fromspox.opset.ai.onnximportv17asopdefgeometric_mean(x:Var,y:Var)->Var:# use the standard Sqrt and Mulreturnop.sqrt(op.mul(x,y))# Create typed model inputs. Each tensor is of rank 1# and has the runtime-determined length 'N'.a=argument(Tensor(float,('N',)))b=argument(Tensor(float,('N',)))# Perform operations on `Var`sc=geometric_mean(a,b)# Build an `onnx.ModelProto` for the given inputs and outputs.model:onnx.ModelProto=build(inputs={'a':a,'b':b},outputs={'c':c})

importnumpyasnpimportonnximportonnx.helperasohfromonnx_array_api.plotting.text_plotimportonnx_simple_text_plotdefmake_euclidean_skl2onnx(input_names:tuple[str]=("X","Y"),output_name:str="Z",elem_type:int=onnx.TensorProto.FLOAT,opset:int|None=None,)->onnx.ModelProto:ifopsetisNone:opset=onnx.defs.onnx_opset_version()fromskl2onnx.algebra.onnx_opsimportOnnxSub,OnnxPow,OnnxReduceSumdxy=OnnxSub(input_names[0],input_names[1],op_version=opset)dxy2=OnnxPow(dxy,np.array([2],dtype=np.int64),op_version=opset)final=OnnxReduceSum(dxy2,op_version=opset,output_names=[output_name])np_type=oh.tensor_dtype_to_np_dtype(elem_type)dummy=np.empty([1],np_type)returnfinal.to_onnx({"X":dummy,"Y":dummy})model=make_euclidean_skl2onnx()print(onnx_simple_text_plot(model))

opset:domain=''version=15input:name='X'type=dtype('float32')shape=['']input:name='Y'type=dtype('float32')shape=['']init:name='Po_Powcst'type=int64shape=(1,)--array([2])Sub(X,Y)->Su_C0Pow(Su_C0,Po_Powcst)->Po_Z0ReduceSum(Po_Z0)->Zoutput:name='Z'type=dtype('float32')shape=[1]

importonnxruntime.training.onnxblockasonnxblockfromonnxruntime.trainingimportartifacts# Define a custom loss block that takes in two inputs# and performs a weighted average of the losses from these# two inputs.classWeightedAverageLoss(onnxblock.Block):def__init__(self):self._loss1=onnxblock.loss.MSELoss()self._loss2=onnxblock.loss.MSELoss()self._w1=onnxblock.blocks.Constant(0.4)self._w2=onnxblock.blocks.Constant(0.6)self._add=onnxblock.blocks.Add()self._mul=onnxblock.blocks.Mul()defbuild(self,loss_input_name1,loss_input_name2):# The build method defines how the block should be stacked on top of# loss_input_name1 and loss_input_name2# Returns weighted average of the two lossesreturnself._add(self._mul(self._w1(),self._loss1(loss_input_name1,target_name="target1")),self._mul(self._w2(),self._loss2(loss_input_name2,target_name="target2")))my_custom_loss=WeightedAverageLoss()# Load the onnx modelmodel_path="model.onnx"base_model=onnx.load(model_path)# Define the parameters that need their gradient computedrequires_grad=["weight1","bias1","weight2","bias2"]frozen_params=["weight3","bias3"]# Now, we can invoke generate_artifacts with this custom loss functionartifacts.generate_artifacts(base_model,requires_grad=requires_grad,frozen_params=frozen_params,loss=my_custom_loss,optimizer=artifacts.OptimType.AdamW)# Successful completion of the above call will generate 4 files in the current working directory,# one for each of the artifacts mentioned above (training_model.onnx, eval_model.onnx, checkpoint, op)

importonnx_graphsurgeonasgsimportnumpyasnpimportonnx# Computes Y = x0 + (a * x1 + b)shape=(1,3,224,224)# Inputsx0=gs.Variable(name="x0",dtype=np.float32,shape=shape)x1=gs.Variable(name="x1",dtype=np.float32,shape=shape)# Intermediate tensorsa=gs.Constant("a",values=np.ones(shape=shape,dtype=np.float32))b=gs.Constant("b",values=np.ones(shape=shape,dtype=np.float32))mul_out=gs.Variable(name="mul_out")add_out=gs.Variable(name="add_out")# OutputsY=gs.Variable(name="Y",dtype=np.float32,shape=shape)nodes=[# mul_out = a * x1gs.Node(op="Mul",inputs=[a,x1],outputs=[mul_out]),# add_out = mul_out + bgs.Node(op="Add",inputs=[mul_out,b],outputs=[add_out]),# Y = x0 + addgs.Node(op="Add",inputs=[x0,add_out],outputs=[Y]),]graph=gs.Graph(nodes=nodes,inputs=[x0,x1],outputs=[Y])onnx.save(gs.export_onnx(graph),"model.onnx")

importonnx_graphsurgeonasgsimportnumpyasnpimportonnxgraph=gs.import_onnx(onnx.load("model.onnx"))# 1. Remove the `b` input of the add nodefirst_add=[nodefornodeingraph.nodesifnode.op=="Add"][0]first_add.inputs=[inpforinpinfirst_add.inputsifinp.name!="b"]# 2. Change the Add to a LeakyRelufirst_add.op="LeakyRelu"first_add.attrs["alpha"]=0.02# 3. Add an identity after the add nodeidentity_out=gs.Variable("identity_out",dtype=np.float32)identity=gs.Node(op="Identity",inputs=first_add.outputs,outputs=[identity_out])graph.nodes.append(identity)# 4. Modify the graph output to be the identity outputgraph.outputs=[identity_out]# 5. Remove unused nodes/tensors, and topologically sort the graph# ONNX requires nodes to be topologically sorted to be considered valid.# Therefore, you should only need to sort the graph when you have added new nodes out-of-order.# In this case, the identity node is already in the correct spot (it is the last node,# and was appended to the end of the list), but to be on the safer side, we can sort anyway.graph.cleanup().toposort()onnx.save(gs.export_onnx(graph),"modified.onnx")

importnumpyasnpfromonnx_array_api.graph_apiimportGraphBuilderfromonnx_array_api.plotting.text_plotimportonnx_simple_text_plotg=GraphBuilder()g.make_tensor_input("X",np.float32,(None,None))g.make_tensor_input("Y",np.float32,(None,None))r1=g.op.Sub("X","Y")r2=g.op.Pow(r1,np.array([2],dtype=np.int64))g.op.ReduceSum(r2,outputs=["Z"])g.make_tensor_output("Z",np.float32,(None,None))onx=g.to_onnx()print(onnx_simple_text_plot(onx))

opset:domain=''version=22input:name='X'type=dtype('float32')shape=['','']input:name='Y'type=dtype('float32')shape=['','']init:name='cst'type=int64shape=(1,)--array([2])Sub(X,Y)->_onx_sub0Pow(_onx_sub0,cst)->_onx_pow0ReduceSum(_onx_pow0)->Zoutput:name='Z'type=dtype('float32')shape=['','']

importnumpyasnpfromonnx_array_api.light_apiimportstartfromonnx_array_api.plotting.text_plotimportonnx_simple_text_plotmodel=(start().vin("X").vin("Y").bring("X","Y").Sub().rename("dxy").cst(np.array([2],dtype=np.int64),"two").bring("dxy","two").Pow().ReduceSum().rename("Z").vout().to_onnx())print(onnx_simple_text_plot(model))

opset:domain=''version=22input:name='X'type=dtype('float32')shape=Noneinput:name='Y'type=dtype('float32')shape=Noneinit:name='two'type=int64shape=(1,)--array([2])Sub(X,Y)->dxyPow(dxy,two)->r1_0ReduceSum(r1_0,keepdims=1,noop_with_empty_axes=0)->Zoutput:name='Z'type=dtype('float32')shape=None

importnumpyasnpfromonnx_array_api.npximportjit_onnxfromonnx_array_api.plotting.text_plotimportonnx_simple_text_plotdefl2_loss(x,y):return((x-y)**2).sum(keepdims=1)jitted_myloss=jit_onnx(l2_loss)dummy=np.array([0],dtype=np.float32)# The function is executed. Only then a onnx graph is created.# One is created depending on the input type.jitted_myloss(dummy,dummy)# get_onnx only works if it was executed once or at least with# the same input typemodel=jitted_myloss.get_onnx()print(onnx_simple_text_plot(model))

opset:domain=''version=18input:name='x0'type=dtype('float32')shape=['']input:name='x1'type=dtype('float32')shape=['']Constant(value=2)->r__1Sub(x0,x1)->r__0CastLike(r__1,r__0)->r__2Pow(r__0,r__2)->r__3ReduceSum(r__3,keepdims=1)->r__4output:name='r__4'type=dtype('float32')shape=[1]