Error: Exporting the operator 'aten::searchsorted' to ONNX opset version 17

[CaioDaumann]

🚀 The feature, motivation and pitch

Dear experts,

I am attempting to export a normalizing flows model built using the Zuko libraries, which are based on PyTorch, to ONNX. However, I am encountering the following error:

Failed to export model: Exporting the operator 'aten::searchsorted' to ONNX opset version 17 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues.

I am following the error suggestion and reporting it as an issue. Also, here is the fully reproducible code:

import torch 
import torch.utils.data as data
import zuko
import numpy as np
import onnxruntime as ort

class WrappedNSF(torch.nn.Module):
    def __init__(self):
        super(WrappedNSF, self).__init__()
        self.flow = zuko.flows.NSF(features=2, transforms=3, hidden_features=(64, 64))
    
    def forward(self, x):
        result = self.flow().transform(x)
        return result

def two_moons(n: int, sigma: float = 1e-1):
    theta = 2 * torch.pi * torch.rand(n)
    label = (theta > torch.pi).float()

    x = torch.stack((
        torch.cos(theta) + label - 1 / 2,
        torch.sin(theta) + label / 2 - 1 / 4,
    ), axis=-1)

    return torch.normal(x, sigma), label

samples, labels = two_moons(16384)

samples_tensor = samples.clone().detach()

trainset = data.TensorDataset(samples, labels)
trainloader = data.DataLoader(trainset, batch_size=64, shuffle=True)

model = WrappedNSF()
model.eval()  

dummy_input = torch.randn(1, 2)  

output = model(dummy_input)
print("Sample output:", output)

# Attempt to export the model to ONNX
try:
    torch.onnx.export(model,               # Wrapped model instance
                    dummy_input,         # Model input (or a tuple for multiple inputs)
                    "wrapped_flow_model.onnx",  # Output ONNX file path
                    export_params=True,    # Store the trained parameter weights inside the model file
                    opset_version=17,      # ONNX version to export the model to
                    do_constant_folding=True,  # Optimization: constant folding
                    input_names=['input'],     # Model's input names
                    output_names=['output'],   # Model's output names
                    dynamic_axes={'input': {0: 'batch_size'},  # Variable length axes
                                    'output': {0: 'batch_size'}})
    print("Model exported successfully.")
except Exception as e:
    print("Failed to export model:", str(e))`

Alternatives

No response

Additional context

torch version: '2.2.1'

[titaiwangms]

https://github.com/microsoft/onnxscript/blob/a5ed07981f1e7c100c5abeb2340a841b9e81e878/onnxscript/function_libs/torch_lib/ops/core.py#L7302-L7312

cc @justinchuby @fatcat-z @shubhambhokare1