The OneHotEncoder operator in ai.onnx.ml was causing segmentation faults when used with custom operators that don't provide complete shape inference information. This occurred because the shape inference function directly accessed the input tensor shape without validating that the input type information was available.

Problem

The crash occurred in the OneHotEncoder shape inference function at this line:

const TensorShapeProto& input_shape = ctx.getInputType(0)->tensor_type().shape();

When a custom operator (like CustomIdentity in the example below) precedes OneHotEncoder and doesn't provide shape information, ctx.getInputType(0) could return nullptr or the input might not have a tensor type, leading to a segmentation fault.

Reproduction

import onnx
from onnx import helper

# Create model with custom operator followed by OneHotEncoder
custom_identity = helper.make_node("CustomIdentity", ["A"], ["B"], domain="com.example")
onehot = helper.make_node("OneHotEncoder", ["B"], ["C"], domain="ai.onnx.ml", cats_strings=["foo", "bar"])

# This would crash with segfault before the fix
onnx.checker.check_model(model, full_check=True)

Solution

Added proper input validation using hasNInputShapes(ctx, 1) before accessing the input shape, following the established pattern used by other operators in the codebase:

// Check if input shape is available before accessing it
if (!hasNInputShapes(ctx, 1)) {
  return;
}

This function safely checks that:

1. The input exists
2. The input type is not null
3. The input has shape information available

Impact

• Prevents segfaults when OneHotEncoder receives inputs from operators without complete type information
• Maintains backward compatibility - normal operation is unchanged when input shapes are available
• Follows established patterns - uses the same defensive coding pattern as ArrayFeatureExtractor and other operators
• Minimal change - only adds 3 lines of validation code without altering core logic

Fixes #7298.

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.