diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp index 428ebffe67..21023debf9 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/DmlOperatorElementWise.cpp @@ -508,18 +508,19 @@ public: Initialize(kernelInfo, std::nullopt, std::nullopt); uint32_t axis = 0; - uint32_t broadcastAxisLength = 0; // If an axis was given explicitly passed (or the default value 1 is set from the schema), // then other inputs are broadcasting to the shape of the input data tensor. if (kernelInfo.HasAttribute(AttrName::Axis, MLOperatorAttributeType::Int)) { + // Avoid validating the axis until later because the axis parameter is ignorable unless + // broadcasting is actually needed. ONNX opset 13 returns a default value of 1 for the + // "axis" attribute even when the attribute doesn't actually exist in the model, which + // would cause a validation failure here. const int32_t signedAxis = gsl::narrow_cast(kernelInfo.GetAttribute(AttrName::Axis)); - axis = Dml::HandleNegativeAxis(signedAxis, outputShapeDimCount); - broadcastAxisLength = outputShape[axis]; + axis = Dml::HandleNegativeAxis(signedAxis, outputShapeDimCount, /*validateAxis*/ false); } - // Explicitly reshape each of the inputs after the first input (scale and zero point tensors). for (uint32_t index = 1, inputCount = gsl::narrow_cast(m_inputTensorDescs.size()); index < inputCount; ++index) { @@ -535,6 +536,8 @@ public: // The 1D vector should have a length equal to the output tensor's dimension on that axis. if (inputTensorShape.size() == 1 && inputTensorShape != outputShape) { + ML_CHECK_VALID_ARGUMENT(axis < outputShapeDimCount); + uint32_t broadcastAxisLength = outputShape[axis]; ML_CHECK_VALID_ARGUMENT(inputTensorShape[0] == broadcastAxisLength); inputTensorShape.insert(inputTensorShape.begin(), axis, 1); inputTensorShape.insert(inputTensorShape.end(), outputShapeDimCount - 1 - axis, 1); diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp index 9169028ae1..7ac27f6598 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.cpp @@ -13,14 +13,14 @@ namespace OperatorHelper // Convert any negative axis into an absolute axis relative to the back end. // So given 3 dimensions, a -1 refers to axis 2, and -3 to axis 0. - uint32_t HandleNegativeAxis(int32_t signedOnnxAxis, uint32_t dimCount) + uint32_t HandleNegativeAxis(int32_t signedOnnxAxis, uint32_t dimCount, bool validateAxis) { if (signedOnnxAxis < 0) { signedOnnxAxis += dimCount; } uint32_t absoluteAxis = gsl::narrow_cast(signedOnnxAxis); - ML_CHECK_VALID_ARGUMENT(absoluteAxis < dimCount); + ML_CHECK_VALID_ARGUMENT(!validateAxis || absoluteAxis < dimCount); return absoluteAxis; } diff --git a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h index 0093160a6a..f060cbdc8b 100644 --- a/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h +++ b/onnxruntime/core/providers/dml/OperatorAuthorHelper/OperatorHelper.h @@ -46,7 +46,7 @@ void FindValueIndices(gsl::span values, T value, /*out*/ std::vector onnxAxes, uint32_t dimCount);