The current implementation is from torchscript exporter symbolic_opset, but there might be a more intuitive way to do this: https://github.com/microsoft/onnxscript/pull/2112#discussion_r2001192021