[Qiming Bao (14H034160212)]

aten_searchsorted is a NotImplementedError stub, so torch.onnx.export(..., dynamo=True) fails on models that call torch.searchsorted, for example top-p sampling (searchsorted(cumsum(softmax(logits)), p)) or quantile binning.

This PR implements aten::searchsorted.Tensor and aten::searchsorted.Scalar. The insertion index is the number of boundaries below each value: boundaries < value for the left side, boundaries <= value for right=True / side="right". It compares values[..., K, 1] against sorted_sequence[..., 1, M] and sums over the last axis, so the same code handles a 1-D sequence (broadcast to every value) and batched N-D sequences. Also handled:

• sorter (applied with GatherElements on the last axis)
• out_int32
• 0-d and Python scalar values

Tests:

• Added searchsorted to ops_test_data.py. All 768 non-empty OpInfo subtests pass. Size-0 samples are skipped with the same fixme as other reductions (ORT reduces size-0 inputs to the wrong shape).
• End-to-end through torch.onnx.export(dynamo=True) + ONNX Runtime, compared with eager: 13/13 cases match. They cover left/right/side, out_int32, batched N-D, sorter, 0-d and scalar values, int64 inputs, a top-p sampling pattern, bucketize, and dynamic batch/sequence/value sizes. On the released onnxscript, every case that reaches searchsorted fails to convert.
• ruff check / ruff format (0.15.8) clean.

[dfedoryshchev]

You missed NaN values here: Less(boundaries, values) and LessOrEqual(boundaries, values) are false for every boundary when the value is NaN, so the count is 0. PyTorch puts NaN at the end: torch.searchsorted(torch.tensor([1., 2., 3.]), torch.tensor([float("nan")])) gives 3, on both sides. The CPU kernel tests !(mid_val >= val) (!(mid_val > val) for right) precisely so that NaN ends up past the last boundary (aten/src/ATen/native/Bucketization.cpp, cus_lower_bound). For float inputs a NaN value needs to map to the size of the last dim, e.g. Or(below, IsNaN(values)) before the ReduceSum.

[Qiming Bao (14H034160212)]

Thanks dfedoryshchev, good catch. Fixed in 1d5d654 as you suggested: for floating-point inputs a NaN value now counts every boundary as below it (Or(below, IsNaN(values)) before the ReduceSum), so it maps to the size of the last dim on both sides. test_searchsorted_places_nan_values_after_every_boundary exports torch.searchsorted with NaN values mixed in, left and right, float32 and float16, and compares with PyTorch; all four cases fail without the change. The OpInfo tests still pass with onnx 1.18 (the CI pin) and 1.23.