Metadata removed when rewrite rule is applied during `torch.onnx.export`

[nickfraser]

Context

Before explaining my issue, please understand (and forgive) the following caveats:

• I'm not sure if my issue is expected behaviour or not. If it is expected behaviour please feel free to close the issue; and
• I'm not interacting with onnxscript directly, but I believe I've traced my issue to a call to onnxscript, so I believe the issue belongs here. Please feel free to correct me if I'm wrong on this.

For further context, please see this issue.

The Issue

With that out-of-the-way, let me explain the issue. It appears when onnxscript does not preserve any metadata when it applies rewrite rules on mergeable operations (in my example, with consecutive transpose operations).

Consider the following code:

import torch

class DoubleTranspose(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, k):
        k = k.transpose(1,2)
        k = k.transpose(-2, -1) # Two transposes get merged and have no metadata
        return 2*k # Some other op which does have metadata

k = torch.rand(1,128,12,32)
model = DoubleTranspose()
output = model(k) # Make sure the forward path runs

with torch.no_grad():
    torch.onnx.export(
        model,
        (k),
        "minimal.onnx",
        input_names=['k'],
        opset_version=18,
        dynamo=True,
        optimize=True, # set `optimize=False` to prevent the transpose nodes from being merged
    )

When optimize=True, the 2 transposes are merged into a single transpose and no longer contain any metadata, as shown:

With optimize=False, you get the two transposes with metadata intact:

In tracing the PyTorch export function, I'm led to the following function in onnxscript when optimize=True during PyTorch export. Is there any way to apply these optimization passes while automagically adding some metadata to the generated node? Either as a argument to torch.onnx.export or as an alternative post-export step would be of interest here.

Environment

• OS: Ubuntu 24.04
• PyTorch version: 2.7.1
• ONNXScript versions: 0.4.0, 0.5.3

[nickfraser]

CC: @jsmonson

[justinchuby]

@titaiwangms @gramalingam

[iksnagreb]

This concrete issue seems to be this replacement pattern: https://github.com/microsoft/onnxscript/blob/v0.5.3/onnxscript/rewriter/rules/common/_basic_rules.py#L276, the replacement creates a new Transpose op ignoring the metadata of the matched transposes. This seems to be the usual way of doing rewrite rules, ignoring the metadata. I see two problems here:

1. The ir.Builder from ONNX IR does not expose the metadata_props, which otherwise should be handled correctly by the ir.Tape, as it seems. This could be added here, similarly to how the domain, version and outputs are already handled: https://github.com/onnx/ir-py/blob/main/src/onnx_ir/_tape.py#L183 (I just toyed around with this, seems to work in priciple, but the builder and tape used here still seem to be the ones from the ir sub-package living within onnxscript, not the standalone ir-py?).
2. More a conceptual issue: Even if you have the mechanism for transferring metadata, how should you decide, in general, which metadata to keep for these types of rules where multiple nodes are fused into one?

[justinchuby]

Even if you have the mechanism for transferring metadata, how should you decide, in general, which metadata to keep for these types of rules where multiple nodes are fused into one?

That's a harder question. I assume we can only merge the known metadata fields 🤔

[gramalingam]

Yes, it would be worth addressing this.

The challenge, as discussed above, is in merging the metadata from multiple nodes. Ideally, it can be done better if we understand the metadata. We should ideally have a compose method defined for a specific metadata that can map a set/list of strings to one string. Ideally, we should standardize metadata within the ONNX standard (with such a composition op defined). And support user-defined or custom metadata by allowing the users to register such a composition op for a specific metadata.

[gramalingam]

Of course, the special case of a single node to single node also shows up and is easy to support and should be.

[nickfraser]

Thanks for the detailed and open discourse on this - it's much appreciated!

Even if you have the mechanism for transferring metadata, how should you decide, in general, which metadata to keep for these types of rules where multiple nodes are fused into one?

Completely agreed that this is non-trivial.

We should ideally have a compose method defined for a specific metadata that can map a set/list of strings to one string. Ideally, we should standardize metadata within the ONNX standard (with such a composition op defined). And support user-defined or custom metadata by allowing the users to register such a composition op for a specific metadata.

My feeling is that this is a great suggestion / compromise. My expectation is that this is relatively "advanced" functionality that is not needed / wanted by most users, when you couple this with the fact that it's obvious how this metadata should be populated (I can imagine many conflicting use-cases) - as such, it seems reasonable to me to put this burden on the user.

@jsmonson, could you comment if you think this suggestion would work for your use-case?

[gramalingam]

A PR that was started on this a while back, just for reference: #2375

[jsmonson]

My feeling is that this is a great suggestion / compromise. My expectation is that this is relatively "advanced" functionality that is not needed / wanted by most users, when you couple this with the fact that it's obvious how this metadata should be populated (I can imagine many conflicting use-cases) - as such, it seems reasonable to me to put this burden on the user.

@jsmonson, could you comment if you think this suggestion would work for your use-case?

Yes. This should work. :)

[gramalingam]

A basic version is in place (from #2682) : currently, we should see some metadata in the output. (Currently, it just retains the metadata from one of the nodes (for each key separately) when multiple nodes are merged. Now, we need to add some more intelligent way of composing multiple ones: while we can simply concatenate, a better solution is possible for source-info that looks like a stack-trace, etc.