OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[ONNX] Improve docstring of onnx symbolic ops (#149668)

Browse Source 
Better examples
Pull Request resolved: https://github.com/pytorch/pytorch/pull/149668
Approved by: https://github.com/titaiwangms
This commit is contained in:
Justin Chu 2025-03-21 01:57:35 +00:00 committed by PyTorch MergeBot
parent 66dd00fca0
commit 362b40939d
1 changed files with 22 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
torch/onnx/ops/__init__.py
View File

@ -82,8 +82,13 @@ def symbolic(
Example:: Example::
class CustomOp(torch.nn.Module): class CustomOp(torch.nn.Module):
def forward(self, x: torch.Tensor): def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.onnx.ops.symbolic( # Normal torch operators can interleave with the symbolic ops during ONNX export
x = x + 1
# Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
# The output tensor will have the specified dtype and shape
val = torch.onnx.ops.symbolic(
"custom_domain::CustomOp", "custom_domain::CustomOp",
(x,), (x,),
dict(attr_key="attr_value"), dict(attr_key="attr_value"),
@ -91,11 +96,12 @@ def symbolic(
shape=x.shape, shape=x.shape,
version=1, version=1,
) )
# This will create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
# The output tensor will have the specified dtype and shape. # The result of the symbolic op can be used in normal torch operations during ONNX export
return torch.nn.functional.relu(val)
# You may then export this model to ONNX using torch.onnx.export. # You may then export this model to ONNX using torch.onnx.export(..., dynamo=True).
Args: Args:
domain_op: The domain and operator name, separated by "::". For example, domain_op: The domain and operator name, separated by "::". For example,
@ -171,8 +177,13 @@ def symbolic_multi_out(
Example:: Example::
class CustomOp(torch.nn.Module): class CustomOp(torch.nn.Module):
def forward(self, x: torch.Tensor): def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.onnx.ops.symbolic( # Normal torch operators can interleave with the symbolic ops during ONNX export
x = x + 1
# Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
# The output tensors will have the specified dtypes and shapes
(out1, out2) = torch.onnx.ops.symbolic(
"custom_domain::CustomOp", "custom_domain::CustomOp",
(x,), (x,),
dict(attr_key="attr_value"), dict(attr_key="attr_value"),
@ -180,11 +191,12 @@ def symbolic_multi_out(
shapes=(x.shape, [1, 2, 3]), shapes=(x.shape, [1, 2, 3]),
version=1, version=1,
) )
# This will create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain.
# The output tensor will have the specified dtype and shape. # The result of the symbolic op can be used in normal torch operations during ONNX export
return torch.nn.functional.relu(out1 + out2)
# You may then export this model to ONNX using torch.onnx.export. # You may then export this model to ONNX using torch.onnx.export(..., dynamo=True).
Args: Args:
domain_op: The domain and operator name, separated by "::". For example, domain_op: The domain and operator name, separated by "::". For example,