@@ -249,9 +249,57 @@ static int64_t EstimateIdentityOutputSizeInBytes(const Node& node) {
249249 return EstimateTensorSizeInBytes (*input_defs[0 ]);
250250}
251251
252+ // ConstantOfShape's output shape is determined by the values of its first input (a shape tensor),
253+ // which is required to be a constant initializer for the node to be eligible for constant folding.
254+ // Compute the output byte size directly from that initializer so we do not have to rely on ONNX
255+ // shape inference having propagated the shape onto the output NodeArg (which is not guaranteed
256+ // for all opsets, all shapes-as-initializers, or all build configurations).
257+ static int64_t EstimateConstantOfShapeOutputSizeInBytes (const Node& node, const Graph& graph) {
258+ const auto & input_defs = node.InputDefs ();
259+ if (input_defs.empty () || input_defs[0 ] == nullptr || !input_defs[0 ]->Exists ()) {
260+ return -1 ;
261+ }
262+
263+ constexpr bool check_outer_scope = true ;
264+ const ONNX_NAMESPACE ::TensorProto* shape_init =
265+ graph.GetConstantInitializer (input_defs[0 ]->Name (), check_outer_scope);
266+ if (shape_init == nullptr ) {
267+ return -1 ;
268+ }
269+
270+ Initializer shape_data{graph, *shape_init, graph.ModelPath ()};
271+ if (shape_data.data_type () != ONNX_NAMESPACE ::TensorProto_DataType_INT64) {
272+ return -1 ;
273+ }
274+
275+ SafeInt<int64_t > num_elements = 1 ;
276+ for (int64_t dim : shape_data.DataAsSpan <int64_t >()) {
277+ if (dim < 0 ) {
278+ return -1 ; // Invalid shape value; let the kernel reject it.
279+ }
280+ num_elements *= dim;
281+ }
282+
283+ // Determine the element size of the output. The ONNX spec for ConstantOfShape defaults the
284+ // element type to float when the optional 'value' attribute is absent.
285+ size_t element_size = sizeof (float );
286+ const auto & attrs = node.GetAttributes ();
287+ auto it = attrs.find (" value"
288+ if (it != attrs.end () && it->second .type () == ONNX_NAMESPACE ::AttributeProto::TENSOR ) {
289+ const auto elem_type = static_cast <ONNX_NAMESPACE ::TensorProto_DataType>(
290+ it->second .t ().data_type ());
291+ const size_t es = GetElementSizeForConstantFolding (elem_type);
292+ if (es != 0 ) {
293+ element_size = es;
294+ }
295+ }
296+
297+ return num_elements * static_cast <int64_t >(element_size);
298+ }
299+
252300// Estimate the total output size in bytes for a node using shape inference results.
253301// Returns -1 if the output size cannot be estimated (e.g., unknown shapes or types).
254- static int64_t EstimateNodeOutputSizeInBytes (const Node& node) {
302+ static int64_t EstimateNodeOutputSizeInBytes (const Node& node, const Graph& graph ) {
255303 if (node.OpType () == " Identity" Domain ().empty ()) {
256304 return EstimateIdentityOutputSizeInBytes (node);
257305 }
@@ -260,6 +308,15 @@ static int64_t EstimateNodeOutputSizeInBytes(const Node& node) {
260308 return EstimateUniqueOutputSizeInBytes (node);
261309 }
262310
311+ if (node.OpType () == " ConstantOfShape" Domain ().empty ()) {
312+ const int64_t size = EstimateConstantOfShapeOutputSizeInBytes (node, graph);
313+ if (size >= 0 ) {
314+ return size;
315+ }
316+ // Fall through to the generic estimator if we could not derive a size from the input
317+ // initializer (e.g., the shape input is not a recognizable constant initializer).
318+ }
319+
263320 SafeInt<int64_t > total_size = 0 ;
264321 for (const auto * output_def : node.OutputDefs ()) {
265322 if (!output_def->Exists ()) {
@@ -391,7 +448,7 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level,
391448 if (max_output_size > 0 ) {
392449 int64_t estimated_size = -1 ;
393450 try {
394- estimated_size = EstimateNodeOutputSizeInBytes (*node);
451+ estimated_size = EstimateNodeOutputSizeInBytes (*node, graph );
395452 } catch (const std::exception&) {
396453 // SafeInt overflow means the size is astronomically large - definitely skip
397454 LOGS (logger, WARNING ) << " Integer overflow while estimating output size of "
0 commit comments