onnx 1.17.0__cp310-cp310-win_amd64.whl → 1.19.0__cp310-cp310-win_amd64.whl

onnx/defs/shape_inference.cc

@@ -6,6 +6,7 @@
 
 #include <vector>
 
+#include "onnx/defs/data_type_utils.h"
 #include "onnx/defs/tensor_proto_util.h"
 
 namespace ONNX_NAMESPACE {
@@ -44,7 +45,7 @@ void propagateElemTypeFromTensorInputToOutput(InferenceContext& ctx, size_t inpu
   }
 }
 
-void propagateElemTypeFromSequenceInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex) {
+static void propagateElemTypeFromSequenceInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex) {
   auto input_type = ctx.getInputType(inputIndex);
   if (nullptr == input_type || input_type->value_case() != TypeProto::kSequenceType) {
     fail_type_inference("Input ", inputIndex, " expected to have sequence type");
@@ -58,7 +59,7 @@ void propagateElemTypeFromSequenceInputToOutput(InferenceContext& ctx, size_t in
   output_type->mutable_sequence_type()->mutable_elem_type()->CopyFrom(input_seq_type.elem_type());
 }
 
-void propagateElemTypeFromOptionalInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex) {
+static void propagateElemTypeFromOptionalInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex) {
   auto input_type = ctx.getInputType(inputIndex);
   if (nullptr == input_type || input_type->value_case() != TypeProto::kOptionalType) {
     fail_type_inference("Input ", inputIndex, " expected to have optional type");
@@ -72,7 +73,7 @@ void propagateElemTypeFromOptionalInputToOutput(InferenceContext& ctx, size_t in
   output_type->mutable_optional_type()->mutable_elem_type()->CopyFrom(input_opt_type.elem_type());
 }
 
-void propagateElemTypeFromMapInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex) {
+static void propagateElemTypeFromMapInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex) {
   auto input_type = ctx.getInputType(inputIndex);
   if (nullptr == input_type || input_type->value_case() != TypeProto::kMapType) {
     fail_type_inference("Input ", inputIndex, " expected to have map type");
@@ -115,7 +116,7 @@ Merge shape information from a source shape into a target shape.
 * prefer target param over source param if mismatched.
 * Fail if there are mismatches in number of dimensions or dimension values.
 */
-void mergeInShapeInfo(const TensorShapeProto& source, TensorShapeProto& target) {
+static void mergeInShapeInfo(const TensorShapeProto& source, TensorShapeProto& target) {
   auto num_source_dims = source.dim_size();
   auto num_target_dims = target.dim_size();
   if (num_source_dims != num_target_dims) {
@@ -183,24 +184,18 @@ void mergeInShapeInfo(const TypeProto_SparseTensor& source, TypeProto_SparseTens
 /// </summary>
 /// <param name="source_shape"></param>
 /// <param name="target_shape">destination shape</param>
-void UnionShapeInfo(const TensorShapeProto& source_shape, TensorShapeProto& target_shape) {
+static void UnionShapeInfo(const TensorShapeProto& source_shape, TensorShapeProto& target_shape) {
   auto source_rank = source_shape.dim_size();
   for (int i = 0; i < source_rank; ++i) {
-    const auto source_dim = source_shape.dim(i);
+    const auto& source_dim = source_shape.dim(i);
     const auto target_dim = target_shape.dim(i);
     bool is_dims_conflict = [&]() {
       if (source_dim.has_dim_value()) {
-        if (target_dim.has_dim_value() && target_dim.dim_value() == source_dim.dim_value()) {
-          return false;
-        }
-        return true;
+        return !target_dim.has_dim_value() || target_dim.dim_value() != source_dim.dim_value();
       }
 
       if (source_dim.has_dim_param()) {
-        if (target_dim.has_dim_param() && target_dim.dim_param() == source_dim.dim_param()) {
-          return false;
-        }
-        return true;
+        return !(target_dim.has_dim_param() && target_dim.dim_param() == source_dim.dim_param());
       }
 
       return (target_dim.has_dim_value() || target_dim.has_dim_param());
@@ -214,7 +209,7 @@ void UnionShapeInfo(const TensorShapeProto& source_shape, TensorShapeProto& targ
 }
 
 template <typename TENSOR_TYPE>
-void UnionShapeInfoForTensor(const TensorShapeProto& source_shape, TENSOR_TYPE& target_type) {
+static void UnionShapeInfoForTensor(const TensorShapeProto& source_shape, TENSOR_TYPE& target_type) {
   if (target_type.has_shape()) {
     TensorShapeProto* target_shape = target_type.mutable_shape();
 
@@ -233,7 +228,7 @@ void UnionShapeInfo(const TensorShapeProto& source_shape, TypeProto_Tensor& targ
   UnionShapeInfoForTensor(source_shape, target_type);
 }
 
-void UnionShapeInfo(const TypeProto_Tensor& source_type, TypeProto_Tensor& target_type) {
+static void UnionShapeInfo(const TypeProto_Tensor& source_type, TypeProto_Tensor& target_type) {
   // The union of a tensor of unknown rank and a tensor of known rank is a tensor of unknown rank.
   // Hence, if the source_type had unknown rank, we clear the shape of the target_type.
   // Otherwise, UnionShapeInfoForTensor handles the rest.
@@ -244,7 +239,7 @@ void UnionShapeInfo(const TypeProto_Tensor& source_type, TypeProto_Tensor& targe
   }
 }
 
-void UnionShapeInfo(const TypeProto_SparseTensor& source_type, TypeProto_SparseTensor& target_type) {
+static void UnionShapeInfo(const TypeProto_SparseTensor& source_type, TypeProto_SparseTensor& target_type) {
   // The union of a tensor of unknown rank and a tensor of known rank is a tensor of unknown rank.
   // Hence, if the source_type had unknown rank, we clear the shape of the target_type.
   // Otherwise, UnionShapeInfoForTensor handles the rest.
@@ -342,7 +337,7 @@ void UnionTypeInfo(const TypeProto& source_type, TypeProto& target_type) {
 // sparse input and outputs dense or vice-versa.
 // If the output value_case is not set, then
 // the input value_case is propagated.
-void propagateTensorElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type) {
+static void propagateTensorElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type) {
   if (nullptr == input_type) {
     fail_type_inference("Input type was null");
   }
@@ -377,7 +372,7 @@ void propagateTensorElemTypeWithValidation(const TypeProto* input_type, TypeProt
   }
 }
 
-void propagateSequenceElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type) {
+static void propagateSequenceElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type) {
   if (nullptr == input_type) {
     fail_type_inference("Input type was null");
   }
@@ -386,7 +381,7 @@ void propagateSequenceElemTypeWithValidation(const TypeProto* input_type, TypePr
     fail_type_inference("Input was expected to have sequence type. Got ", input_type->value_case());
   }
 
-  auto input_seq_type = input_type->sequence_type();
+  const auto& input_seq_type = input_type->sequence_type();
 
   if (input_seq_type.has_elem_type()) {
     propagateElemTypeWithValidation(
@@ -396,7 +391,7 @@ void propagateSequenceElemTypeWithValidation(const TypeProto* input_type, TypePr
   }
 }
 
-void propagateOptionalElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type) {
+static void propagateOptionalElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type) {
   if (nullptr == input_type) {
     fail_type_inference("Input type was null");
   }
@@ -405,7 +400,7 @@ void propagateOptionalElemTypeWithValidation(const TypeProto* input_type, TypePr
     fail_type_inference("Input was expected to have optional type. Got ", input_type->value_case());
   }
 
-  auto input_opt_type = input_type->optional_type();
+  const auto& input_opt_type = input_type->optional_type();
 
   if (input_opt_type.has_elem_type()) {
     propagateElemTypeWithValidation(
@@ -415,7 +410,7 @@ void propagateOptionalElemTypeWithValidation(const TypeProto* input_type, TypePr
   }
 }
 
-void propagateMapElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type) {
+static void propagateMapElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type) {
   if (nullptr == input_type) {
     fail_type_inference("Input type was null");
   }
@@ -424,7 +419,7 @@ void propagateMapElemTypeWithValidation(const TypeProto* input_type, TypeProto*
     fail_type_inference("Input was expected to have map type. Got ", input_type->value_case());
   }
 
-  auto input_map_type = input_type->map_type();
+  const auto& input_map_type = input_type->map_type();
 
   if (!input_map_type.has_key_type()) {
     fail_type_inference("Key type of map input was unknown");
@@ -459,29 +454,34 @@ void propagateElemTypeWithValidation(const TypeProto* input_type, TypeProto* out
 }
 
 TensorShapeProto getShapeInput(const InferenceContext& ctx, size_t input_index, bool& found) {
+  return getShapeInput(ctx, input_index, false, found);
+}
+
+TensorShapeProto
+getShapeInput(const InferenceContext& ctx, size_t input_index, bool fail_if_negative_value, bool& found) {
   TensorShapeProto shape_input;
 
+  found = false;
+
   // First, check initializer.
   const TensorProto* shape_initializer = ctx.getInputData(input_index);
   if (shape_initializer) {
-    const std::vector<int64_t>& shape_data = ParseData<int64_t>(shape_initializer);
+    const std::vector<int64_t> shape_data = ParseData<int64_t>(shape_initializer);
     for (const int64_t& e : shape_data) {
       shape_input.add_dim()->set_dim_value(e);
     }
     found = true;
-    return shape_input;
   }
 
   // Then, check symbolic input.
   const TensorShapeProto* symbolic_input = ctx.getSymbolicInput(input_index);
-  if (symbolic_input) {
+  if (!found && symbolic_input) {
     shape_input.CopyFrom(*symbolic_input);
     found = true;
-    return shape_input;
   }
 
   // Try rank inference.
-  if (hasInputShape(ctx, input_index)) {
+  if (!found && hasInputShape(ctx, input_index)) {
     const TensorShapeProto& shape_input_shape = getInputShape(ctx, input_index);
     if (shape_input_shape.dim_size() != 1) {
       fail_shape_inference("shape input must be 1D tensor");
@@ -493,12 +493,19 @@ TensorShapeProto getShapeInput(const InferenceContext& ctx, size_t input_index,
         shape_input.add_dim();
       }
       found = true;
-      return shape_input;
     }
   }
 
-  // Shape input was not found.
-  found = false;
+  if (found && fail_if_negative_value) {
+    int dims_size = shape_input.dim_size();
+    for (int i = 0; i < dims_size; ++i) {
+      const auto& dim = shape_input.dim(i);
+      if (dim.has_dim_value() && dim.dim_value() < 0) {
+        fail_shape_inference("shape input tensor must have non-negative elements");
+      }
+    }
+  }
+
   return shape_input;
 }
 

onnx/defs/shape_inference.h

@@ -10,7 +10,7 @@
 #include <utility>
 #include <vector>
 
-#include "onnx/defs/data_type_utils.h"
+#include "onnx/common/common.h"
 #include "onnx/proto_utils.h"
 #include "onnx/string_utils.h"
 
@@ -28,8 +28,8 @@ struct ShapeInferenceOptions {
   // Enables data propagation for limited operators
   // to perform shape computation
   bool enable_data_propagation;
-  ShapeInferenceOptions(bool check_type_val = false, int strict_mode_val = 0, bool data_prop_val = false)
-      : check_type(check_type_val), error_mode(strict_mode_val), enable_data_propagation(data_prop_val){};
+  explicit ShapeInferenceOptions(bool check_type_val = false, int strict_mode_val = 0, bool data_prop_val = false)
+      : check_type(check_type_val), error_mode(strict_mode_val), enable_data_propagation(data_prop_val) {}
 };
 
 // Maintains a SymbolTable for symbolic shape inference
@@ -61,7 +61,7 @@ class InferenceError final : public std::runtime_error {
  public:
   using std::runtime_error::runtime_error;
 
-  InferenceError(const std::string& message) : std::runtime_error(message) {}
+  explicit InferenceError(const std::string& message) : std::runtime_error(message) {}
 
   const char* what() const noexcept override {
     if (!expanded_message_.empty()) {
@@ -100,8 +100,11 @@ struct InferenceContext {
   virtual const TensorProto* getInputData(size_t index) const = 0;
   virtual size_t getNumOutputs() const = 0;
   virtual TypeProto* getOutputType(size_t index) = 0;
+  virtual bool hasOutput(size_t index) {
+    return (index < getNumOutputs() && (getOutputType(index) != nullptr));
+  }
   virtual GraphInferencer* getGraphAttributeInferencer(const std::string& attribute_name) = 0;
-  virtual ~InferenceContext() {}
+  virtual ~InferenceContext() = default;
   virtual const SparseTensorProto* getInputSparseData(size_t index) const = 0;
   // Gets the shape inputs computed by partial data propagation.
   virtual const TensorShapeProto* getSymbolicInput(size_t index) const = 0;
@@ -128,7 +131,7 @@ struct DataPropagationContext {
   virtual const TypeProto* getInputType(size_t index) const = 0;
   virtual size_t getNumOutputs() const = 0;
   virtual const TypeProto* getOutputType(size_t index) const = 0;
-  virtual ~DataPropagationContext() {}
+  virtual ~DataPropagationContext() = default;
   virtual const TensorShapeProto* getInputData(size_t index) = 0;
   virtual void addOutputData(size_t index, TensorShapeProto&& tp) = 0;
 };
@@ -138,13 +141,13 @@ using DataPropagationFunction = std::function<void(DataPropagationContext&)>;
 
 // This no-op inference function is used for operators without an
 // inference implementation.
-inline void dummyInferenceFunction(InferenceContext&){};
+inline void dummyInferenceFunction(InferenceContext&) {}
 
 // This no-op data propagation function is used for operators without a defined data propagator
-inline void dummyDataPropagationFunction(DataPropagationContext&){};
+inline void dummyDataPropagationFunction(DataPropagationContext&) {}
 
 template <typename T>
-inline bool getRepeatedAttribute(InferenceContext& ctx, std::string attr_name, std::vector<T>& values) {
+inline bool getRepeatedAttribute(InferenceContext& ctx, const std::string& attr_name, std::vector<T>& values) {
   const auto* attr = ctx.getAttribute(attr_name);
   if (attr) {
     values = RetrieveValues<T>(*attr);
@@ -154,14 +157,14 @@ inline bool getRepeatedAttribute(InferenceContext& ctx, std::string attr_name, s
   }
 }
 
-inline int64_t getAttribute(InferenceContext& ctx, const std::string& attributeName, int64_t defaultValue) {
+inline int64_t getAttribute(const InferenceContext& ctx, const std::string& attributeName, int64_t defaultValue) {
   auto attr_proto = ctx.getAttribute(attributeName);
   if ((nullptr != attr_proto) && attr_proto->has_i())
     return attr_proto->i();
   return defaultValue;
 }
 
-inline int64_t getAttribute(DataPropagationContext& ctx, const std::string& attributeName, int64_t defaultValue) {
+inline int64_t getAttribute(const DataPropagationContext& ctx, const std::string& attributeName, int64_t defaultValue) {
   auto attr_proto = ctx.getAttribute(attributeName);
   if ((nullptr != attr_proto) && attr_proto->has_i())
     return attr_proto->i();
@@ -169,14 +172,16 @@ inline int64_t getAttribute(DataPropagationContext& ctx, const std::string& attr
 }
 
 inline std::string
-getAttribute(InferenceContext& ctx, const std::string& attributeName, const std::string& defaultValue) {
+getAttribute(const InferenceContext& ctx, const std::string& attributeName, const std::string& defaultValue) {
   auto attr_proto = ctx.getAttribute(attributeName);
   if ((nullptr != attr_proto) && attr_proto->has_s())
     return attr_proto->s();
   return defaultValue;
 }
 
-inline TensorShapeProto::Dimension operator*(TensorShapeProto::Dimension dim1, TensorShapeProto::Dimension dim2) {
+inline TensorShapeProto::Dimension operator*(
+    const TensorShapeProto::Dimension& dim1,
+    const TensorShapeProto::Dimension& dim2) {
   TensorShapeProto::Dimension result;
   if (dim1.has_dim_value() && dim2.has_dim_value()) {
     result.set_dim_value(dim1.dim_value() * dim2.dim_value());
@@ -197,7 +202,7 @@ std::pair<int, int> getAttributeElementTypeAndLength(
     const InferenceContext& ctx,
     const std::initializer_list<std::string>& attribute_names);
 
-inline TensorShapeProto::Dimension operator*(TensorShapeProto::Dimension dim1, int64_t dim2) {
+inline TensorShapeProto::Dimension operator*(const TensorShapeProto::Dimension& dim1, int64_t dim2) {
   TensorShapeProto::Dimension result;
   if (dim1.has_dim_value()) {
     result.set_dim_value(dim1.dim_value() * dim2);
@@ -207,7 +212,7 @@ inline TensorShapeProto::Dimension operator*(TensorShapeProto::Dimension dim1, i
   return result;
 }
 
-inline TensorShapeProto::Dimension operator/(TensorShapeProto::Dimension dim1, int64_t dim2) {
+inline TensorShapeProto::Dimension operator/(const TensorShapeProto::Dimension& dim1, int64_t dim2) {
   TensorShapeProto::Dimension result;
   if (dim1.has_dim_value()) {
     result.set_dim_value(dim1.dim_value() / dim2);
@@ -248,11 +253,11 @@ inline void setTensorElementType(int32_t elem_type, TypeProto::ValueCase value_c
   }
 }
 
-void propagateElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type);
+ONNX_API void propagateElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type);
 
-void propagateElemTypeFromInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
+ONNX_API void propagateElemTypeFromInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
 
-void propagateElemTypeFromTensorInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
+ONNX_API void propagateElemTypeFromTensorInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
 
 inline void propagateElemTypeFromDtypeToOutput(
     InferenceContext& ctx,
@@ -350,7 +355,7 @@ inline const TensorShapeProto& getInputShape(const InferenceContext& ctx, size_t
   }
 }
 
-inline const TensorShapeProto* getOptionalInputShape(InferenceContext& ctx, size_t n) {
+inline const TensorShapeProto* getOptionalInputShape(const InferenceContext& ctx, size_t n) {
   const auto* input_type = ctx.getInputType(n);
 
   if (input_type == nullptr) {
@@ -446,7 +451,7 @@ inline void propagateShapeFromInputToOutput(InferenceContext& ctx, size_t inputI
   propagateShape(input_type, output_type);
 }
 
-inline void propagateShapeAndTypeFromFirstInput(InferenceContext& ctx) {
+ONNX_API inline void propagateShapeAndTypeFromFirstInput(InferenceContext& ctx) {
   propagateElemTypeFromInputToOutput(ctx, 0, 0);
   if (!hasNInputShapes(ctx, 1)) {
     return;
@@ -517,9 +522,7 @@ inline void propagateElemTypeFromAttributeToOutput(
 }
 
 inline TensorShapeProto* getTensorMutableShape(TypeProto::ValueCase value_case, TypeProto& type) {
-  if (value_case == TypeProto::kTensorType) {
-    return type.mutable_tensor_type()->mutable_shape();
-  } else if (value_case == TypeProto::kSparseTensorType) {
+  if (value_case == TypeProto::kTensorType || value_case == TypeProto::kSparseTensorType) {
     return type.mutable_tensor_type()->mutable_shape();
   }
   return nullptr;
@@ -533,9 +536,17 @@ getOutputShape(InferenceContext& ctx, size_t n, TypeProto::ValueCase default_typ
   }
   const auto output_value_case = output_type->value_case();
   if (output_value_case == TypeProto::kTensorType || output_value_case == TypeProto::kSparseTensorType) {
-    return getTensorMutableShape(output_value_case, *output_type);
+    auto output_shape = getTensorMutableShape(output_value_case, *output_type);
+    if (output_shape == nullptr) {
+      fail_type_inference("Output ", n, " expected to have tensor or sparse type in ", ctx.getDisplayName(), ".");
+    }
+    return output_shape;
   } else if (output_value_case == TypeProto::VALUE_NOT_SET) {
-    return getTensorMutableShape(default_type, *output_type);
+    auto output_shape = getTensorMutableShape(default_type, *output_type);
+    if (output_shape == nullptr) {
+      fail_type_inference("Output ", n, " expected to have tensor or sparse type in ", ctx.getDisplayName(), ".");
+    }
+    return output_shape;
   } else {
     fail_type_inference("Output ", n, " expected to have tensor type in ", ctx.getDisplayName(), ".");
   }
@@ -584,6 +595,11 @@ inline void updateOutputShape(
 // Otherwise, `false` is stored, which means that returned TensorShapeProto does not make sense.
 TensorShapeProto getShapeInput(const InferenceContext& ctx, size_t input_index, bool& found);
 
+// Argument `fail_if_negative_value` is used to control whether negative values are allowed in the shape. The shape
+// check would fail if not.
+TensorShapeProto
+getShapeInput(const InferenceContext& ctx, size_t input_index, bool fail_if_negative_value, bool& found);
+
 // Infer shape of an output from the value of a specified attribute, which is
 // expected to be a list of integers specifying a valid shape.
 inline void propagateShapeFromAttributeToOutput(
@@ -613,9 +629,9 @@ inline void multidirectionalBroadcastShapeInference(
     TensorShapeProto& resultShape) {
   int result_shape_size = 0;
   // Get the result shape size.
-  for (size_t i = 0; i < shapes.size(); ++i) {
-    if (shapes[i]->dim_size() > result_shape_size) {
-      result_shape_size = shapes[i]->dim_size();
+  for (auto shape : shapes) {
+    if (shape->dim_size() > result_shape_size) {
+      result_shape_size = shape->dim_size();
     }
   }
 
@@ -623,13 +639,13 @@ inline void multidirectionalBroadcastShapeInference(
     int64_t dim_value = 1;
     TensorShapeProto_Dimension symbolic_dim;
     int num_symbolic_dims = 0;
-    for (size_t j = 0; j < shapes.size(); ++j) {
-      if (i < result_shape_size - shapes[j]->dim_size()) {
+    for (auto shape : shapes) {
+      if (i < result_shape_size - shape->dim_size()) {
         // Shape j will be filled with 1 at dimension i;
         continue;
       }
 
-      auto dim_i_j = shapes[j]->dim(i - result_shape_size + shapes[j]->dim_size());
+      auto dim_i_j = shape->dim(i - result_shape_size + shape->dim_size());
       if (dim_i_j.has_dim_value()) {
         if (dim_i_j.dim_value() != 1) {
           if (dim_value != dim_i_j.dim_value() && dim_value != 1) {
@@ -704,9 +720,8 @@ inline void mergeInDimensionInfo(
     } else {
       target_dim.set_dim_value(source_value);
     }
-  } else if (target_dim.has_dim_value()) {
+  } else if (target_dim.has_dim_value() || target_dim.has_dim_param()) {
     // if target has a value we preserve it so do nothing
-  } else if (target_dim.has_dim_param()) {
     // prefer target param over source
   } else if (source_dim.has_dim_param()) {
     target_dim.set_dim_param(source_dim.dim_param());
@@ -772,7 +787,7 @@ static constexpr T narrow_cast(U&& u) noexcept {
   return static_cast<T>(std::forward<U>(u));
 }
 
-inline void checkInputRank(InferenceContext& ctx, size_t input_index, int expected_rank) {
+inline void checkInputRank(const InferenceContext& ctx, size_t input_index, int expected_rank) {
   // We check the rank only if a rank is known for the input:
   if (hasInputShape(ctx, input_index)) {
     auto rank = getInputShape(ctx, input_index).dim_size();
@@ -821,10 +836,9 @@ inline void unifyDim(const Dim& source_dim, Dim& target_dim) {
     } else {
       target_dim.set_dim_value(source_value);
     }
-  } else if (target_dim.has_dim_value()) {
+  } else if (target_dim.has_dim_value() || target_dim.has_dim_param()) {
     // if target has a value we preserve it.
     // we cannot set source dim value.
-  } else if (target_dim.has_dim_param()) {
     // prefer target param over source
     // we cannot currently unify the dim_params
   } else if (source_dim.has_dim_param()) {
@@ -832,7 +846,7 @@ inline void unifyDim(const Dim& source_dim, Dim& target_dim) {
   }
 }
 
-inline void unifyInputDim(InferenceContext& ctx, size_t input_index, int dim_index, Dim& dim) {
+inline void unifyInputDim(const InferenceContext& ctx, size_t input_index, int dim_index, Dim& dim) {
   // We unify the dimensions only if it is available for specified input:
   if (hasInputShape(ctx, input_index)) {
     auto& input_shape = getInputShape(ctx, input_index);
@@ -908,13 +922,26 @@ void checkAxesRange(Axes& axes, int rank) {
 // checkDuplicateAxes: Check that there are no duplicated axes
 template <typename Axes>
 void checkDuplicateAxes(Axes& axes, int rank) {
-  std::vector<bool> tmp(rank, false);
+  std::vector<bool> tmp(static_cast<size_t>(rank), false);
   for (auto axis : axes) {
-    int actual_axis = axis < 0 ? axis + rank : axis;
+    size_t actual_axis = static_cast<size_t>(axis < 0 ? axis + rank : axis);
     if (tmp[actual_axis])
       fail_shape_inference("Axis ", axis, " is referred to more than once.");
     tmp[actual_axis] = true;
   }
 }
 
+// Shape inference functions for various ONNX operators.
+// Users can use these functions to implement shape inference for custom operators
+// by calling them in their own inference functions.
+ONNX_API void RNNShapeInference(InferenceContext& ctx);
+ONNX_API void convPoolShapeInference(
+    InferenceContext& ctx,
+    bool use_dilation,
+    bool require_kernel_shape,
+    int input1Idx,
+    int input2Idx);
+ONNX_API void convTransposeShapeInference(InferenceContext& ctx);
+ONNX_API void globalPoolTypeShapeInference(InferenceContext& ctx);
+
 } // namespace ONNX_NAMESPACE