// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

#include "onnx_extensions/shapes/dispatch_table.h"

#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>

#include "onnx_core/shapes/dispatch_table.h"
#include "onnx_proto/onnx_helper.h"

#include "onnx_extensions/shapes/shapes/controlflow/shape_controlflow.h"
#include "onnx_extensions/shapes/shapes/generator/shape_generator.h"
#include "onnx_extensions/shapes/shapes/image/shape_image.h"
#include "onnx_extensions/shapes/shapes/logical/shape_logical.h"
#include "onnx_extensions/shapes/shapes/math/shape_math.h"
#include "onnx_extensions/shapes/shapes/nn/shape_nn.h"
#include "onnx_extensions/shapes/shapes/optional/shape_optional.h"
#include "onnx_extensions/shapes/shapes/preview/shape_preview.h"
#include "onnx_extensions/shapes/shapes/quantization/shape_quantization.h"
#include "onnx_extensions/shapes/shapes/reduction/shape_reduction.h"
#include "onnx_extensions/shapes/shapes/rt/shape_rt.h"
#include "onnx_extensions/shapes/shapes/sequence/shape_sequence.h"
#include "onnx_extensions/shapes/shapes/tensor/shape_tensor.h"
#include "onnx_extensions/shapes/shapes/text/shape_text.h"
#include "onnx_extensions/shapes/shapes/traditionalml/shape_traditionalml.h"
#include "onnx_extensions/shapes/shapes/training/shape_training.h"

namespace ONNX_LIGHT_NAMESPACE::onnx_shapes {
namespace shapes {

namespace {

// Verifies the node declares at least `expected` inputs.
void RequireInputs(const NodeProto &node, int expected) {
  EXT_ENFORCE_INVALID(node.input_size() >= expected, "ComputeShapeNode: op '", node.op_type(),
                      "' expects at least ", std::to_string(expected), " input(s), got ",
                      std::to_string(node.input_size()), ".");
}

using ::onnx_light::core::shapes::ComputeShapeFn;

// Built-in table of every ``onnx_shapes`` shape function, keyed by
// ``"<domain>:<op_type>"``. Only used to populate the shared
// ``core::shapes`` dispatch table via :cpp:func:`RegisterShapeFunctions`;
// never queried directly by shape inference.
const std::unordered_map<std::string, ComputeShapeFn> &BuiltinShapeFunctions() {
  static const std::unordered_map<std::string, ComputeShapeFn> table = {
      {"ai.onnx:Abs",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeAbs(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Acos",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeAcos(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Acosh",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeAcosh(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Add",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeAdd(ctx, node, node.input(0).c_str(),
                               node.input(1).c_str());
       }},
      {"ai.onnx:AffineGrid",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeAffineGrid(ctx, node);
       }},
      {"ai.onnx:And",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeAnd(ctx, node, node.input(0).c_str(),
                                  node.input(1).c_str());
       }},
      {"ai.onnx:ArgMax",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         reduction::ComputeShapeArgReduce(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:ArgMin",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         reduction::ComputeShapeArgReduce(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Asin",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeAsin(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Asinh",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeAsinh(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Atan",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeAtan(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Atanh",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeAtanh(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Attention",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         const std::string& q_name = node.input(0);
         const std::string& k_name = node.input(1);
         const std::string& v_name = node.input(2);
         const std::string& past_k_name = node.input_size() > 4 ? node.input(4) : utils::String::empty_string();
         const std::string& past_v_name = node.input_size() > 5 ? node.input(5) : utils::String::empty_string();
         nn::ComputeShapeAttention(ctx, node, q_name.c_str(), k_name.c_str(), v_name.c_str(),
                                   past_k_name.empty() ? nullptr : past_k_name.c_str(),
                                   past_v_name.empty() ? nullptr : past_v_name.c_str());
       }},
      {"ai.onnx:AveragePool",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeAveragePool(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:BatchNormalization",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 5);
         const std::string x_name = node.input(0);
         const std::string mean_name = node.input(3);
         nn::ComputeShapeBatchNormalization(ctx, node, x_name.c_str(),
                                            mean_name.empty() ? nullptr : mean_name.c_str());
       }},
      {"ai.onnx:Dropout",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        const std::string& data_name = node.input(0);
        const std::string& ratio_name =
            node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
        const std::string& training_mode_name =
            node.input_size() >= 3 ? node.input(2) : utils::String::empty_string();
        nn::ComputeShapeDropout(ctx, node, data_name.c_str(),
                                ratio_name.empty() ? nullptr : ratio_name.c_str(),
                                training_mode_name.empty() ? nullptr : training_mode_name.c_str());
       }},
      {"ai.onnx:EyeLike",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeEyeLike(ctx, node);
       }},
      {"ai.onnx:Range",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         generator::ComputeShapeRange(ctx, node);
       }},
      {"ai.onnx:Flatten",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeFlatten(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:GroupNormalization",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         nn::ComputeShapeGroupNormalization(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:InstanceNormalization",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         nn::ComputeShapeInstanceNormalization(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:LayerNormalization",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         nn::ComputeShapeLayerNormalization(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:LinearAttention",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         const std::string& q_name = node.input(0);
         const std::string& k_name = node.input(1);
         const std::string& v_name = node.input(2);
         const std::string& past_state_name =
             node.input_size() > 3 ? node.input(3) : utils::String::empty_string();
         nn::ComputeShapeLinearAttention(
             ctx, node, q_name.c_str(), k_name.c_str(), v_name.c_str(),
             past_state_name.empty() ? nullptr : past_state_name.c_str());
       }},
      {"ai.onnx:MeanVarianceNormalization",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeMeanVarianceNormalization(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:RMSNormalization",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         nn::ComputeShapeRMSNormalization(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Bernoulli",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeBernoulli(ctx, node);
       }},
      {"ai.onnx:RandomNormal",
       [](ShapesContext &ctx, const NodeProto &node) {
         generator::ComputeShapeRandomNormal(ctx, node);
       }},
      {"ai.onnx:RandomNormalLike",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeRandomNormalLike(ctx, node);
       }},
      {"ai.onnx:RandomUniform",
       [](ShapesContext &ctx, const NodeProto &node) {
         generator::ComputeShapeRandomUniform(ctx, node);
       }},
      {"ai.onnx:RandomUniformLike",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeRandomUniformLike(ctx, node);
       }},
      {"ai.onnx:BitCast",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeBitCast(ctx, node);
       }},
      {"ai.onnx:BitShift",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeBitShift(ctx, node, node.input(0).c_str(),
                                       node.input(1).c_str());
       }},
      {"ai.onnx:BitwiseAnd",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeBitwiseAnd(ctx, node, node.input(0).c_str(),
                                         node.input(1).c_str());
       }},
      {"ai.onnx:BitwiseNot",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         logical::ComputeShapeBitwiseNot(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:BitwiseOr",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeBitwiseOr(ctx, node, node.input(0).c_str(),
                                        node.input(1).c_str());
       }},
      {"ai.onnx:BitwiseXor",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeBitwiseXor(ctx, node, node.input(0).c_str(),
                                         node.input(1).c_str());
       }},
      {"ai.onnx:BlackmanWindow",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeBlackmanWindow(ctx, node);
       }},
      {"ai.onnx:Cast",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeCast(ctx, node);
       }},
      {"ai.onnx:CastLike",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeCastLike(ctx, node);
       }},
      {"ai.onnx:Compress",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeCompress(ctx, node);
       }},
      {"ai.onnx:Constant",
       [](ShapesContext &ctx, const NodeProto &node) {
         generator::ComputeShapeConstant(ctx, node);
       }},
      {"ai.onnx:ConstantOfShape",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeConstantOfShape(ctx, node);
       }},
      {"ai.onnx:Concat",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeConcat(ctx, node);
       }},
      {"ai.onnx:Cos",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeCos(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Cosh",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeCosh(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:CumProd",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeCumProd(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:CumSum",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeCumSum(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:DFT",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeDFT(ctx, node);
       }},
      {"ai.onnx:Ceil",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeCeil(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Clip",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeClip(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Floor",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeFloor(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Reciprocal",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeReciprocal(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Round",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeRound(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Sign",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSign(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Sin",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSin(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Sinh",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSinh(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Sqrt",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSqrt(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:STFT",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeSTFT(ctx, node);
       }},
      {"ai.onnx:Sum",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSum(ctx, node);
       }},
      {"ai.onnx:Swish",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSwish(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:SwiGLU",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeSwiGLU(ctx, node, node.input(0).c_str(), node.input(1).c_str());
       }},
      {"ai.onnx:Tan",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeTan(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Tanh",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeTanh(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:ThresholdedRelu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeThresholdedRelu(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Relu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeRelu(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:LeakyRelu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeLeakyRelu(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Elu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeElu(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Celu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeCelu(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Gelu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeGelu(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Selu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSelu(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:TopK",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeTopK(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:DequantizeLinear",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         const std::string x_name = node.input(0);
         const std::string x_scale_name = node.input(1);
         quantization::ComputeShapeDequantizeLinear(ctx, node, x_name.c_str(),
                                                    x_scale_name.c_str());
       }},
      {"ai.onnx:Col2Im",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         nn::ComputeShapeCol2Im(ctx, node, node.input(0).c_str(),
                                node.input(1).c_str(),
                                node.input(2).c_str());
       }},
      {"ai.onnx:Conv",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         nn::ComputeShapeConv(ctx, node, node.input(0).c_str(),
                              node.input(1).c_str());
       }},
      {"ai.onnx:ConvInteger",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         nn::ComputeShapeConvInteger(ctx, node, node.input(0).c_str(),
                                     node.input(1).c_str());
       }},
      {"ai.onnx:ConvTranspose",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         nn::ComputeShapeConvTranspose(ctx, node, node.input(0).c_str(),
                                       node.input(1).c_str());
       }},
      {"ai.onnx:DeformConv",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         nn::ComputeShapeDeformConv(ctx, node, node.input(0).c_str(),
                                    node.input(1).c_str());
       }},
      {"ai.onnx:Det",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeDet(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Div",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeDiv(ctx, node, node.input(0).c_str(),
                               node.input(1).c_str());
       }},
      {"ai.onnx:Einsum",
      [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        math::ComputeShapeEinsum(ctx, node);
      }},
      {"ai.onnx:Erf",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
        math::ComputeShapeErf(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Exp",
      [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        math::ComputeShapeExp(ctx, node, node.input(0).c_str());
      }},
      {"ai.onnx:Equal",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeEqual(ctx, node, node.input(0).c_str(),
                                    node.input(1).c_str());
       }},
      {"ai.onnx:Where",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         logical::ComputeShapeWhere(ctx, node, node.input(0).c_str(),
                                    node.input(1).c_str(),
                                    node.input(2).c_str());
       }},
      {"ai.onnx:Or",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeOr(ctx, node, node.input(0).c_str(),
                                 node.input(1).c_str());
       }},
      {"ai.onnx:Xor",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeXor(ctx, node, node.input(0).c_str(),
                                  node.input(1).c_str());
       }},
      {"ai.onnx:Expand",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeExpand(ctx, node);
       }},
      {"ai.onnx:Gemm",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeGemm(ctx, node, node.input(0).c_str(),
                                node.input(1).c_str());
       }},
      {"ai.onnx:GlobalAveragePool",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeGlobalPool(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:GlobalLpPool",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeGlobalPool(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:GlobalMaxPool",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeGlobalPool(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Greater",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeGreater(ctx, node, node.input(0).c_str(),
                                      node.input(1).c_str());
       }},
      {"ai.onnx:GreaterOrEqual",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeGreaterOrEqual(ctx, node, node.input(0).c_str(),
                                             node.input(1).c_str());
       }},
      {"ai.onnx:GridSample",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeGridSample(ctx, node);
       }},
      {"ai.onnx:Gather",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeGather(ctx, node);
       }},
      {"ai.onnx:GatherElements",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeGatherElements(ctx, node);
       }},
      {"ai.onnx:GatherND",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeGatherND(ctx, node);
       }},
      {"ai.onnx:TensorScatter",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeTensorScatter(ctx, node);
       }},
      {"ai.onnx:Scatter",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         tensor::ComputeShapeScatter(ctx, node);
       }},
      {"ai.onnx:ScatterElements",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         tensor::ComputeShapeScatterElements(ctx, node);
       }},
      {"ai.onnx:ScatterND",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         tensor::ComputeShapeScatterND(ctx, node);
       }},
      {"ai.onnx:HammingWindow",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeHammingWindow(ctx, node);
       }},
      {"ai.onnx:HannWindow",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeHannWindow(ctx, node);
       }},
      {"ai.onnx:HardSigmoid",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeHardSigmoid(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:HardSwish",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeHardSwish(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Hardmax",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeHardmax(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:If",
       [](ShapesContext &ctx, const NodeProto &node) { controlflow::ComputeShapeIf(ctx, node); }},
      {"ai.onnx:ImageDecoder",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         image::ComputeShapeImageDecoder(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:IsInf",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         logical::ComputeShapeIsInf(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:IsNaN",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         logical::ComputeShapeIsNaN(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Loop",
       [](ShapesContext &ctx, const NodeProto &node) {
         controlflow::ComputeShapeLoop(ctx, node);
       }},
      {"ai.onnx:Scan",
       [](ShapesContext &ctx, const NodeProto &node) {
         controlflow::ComputeShapeScan(ctx, node);
       }},
      {"ai.onnx:Log",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeLog(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Less",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeLess(ctx, node, node.input(0).c_str(),
                                   node.input(1).c_str());
       }},
      {"ai.onnx:LessOrEqual",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         logical::ComputeShapeLessOrEqual(ctx, node, node.input(0).c_str(),
                                          node.input(1).c_str());
       }},
      {"ai.onnx:LRN",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeLRN(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:LpNormalization",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeLpNormalization(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:LpPool",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeLpPool(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:MaxPool",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         nn::ComputeShapeMaxPool(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:MaxUnpool",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         const char *output_shape = nullptr;
         std::string output_shape_name;
         if (node.input_size() >= 3 && !node.input(2).empty()) {
           output_shape_name = node.input(2);
           output_shape = output_shape_name.c_str();
         }
         nn::ComputeShapeMaxUnpool(ctx, node, node.input(0).c_str(),
                                   node.input(1).c_str(), output_shape);
       }},
      {"ai.onnx:MaxRoiPool",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         nn::ComputeShapeMaxRoiPool(ctx, node, node.input(0).c_str(),
                                    node.input(1).c_str());
       }},
      {"ai.onnx:Mul",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeMul(ctx, node, node.input(0).c_str(),
                               node.input(1).c_str());
       }},
      {"ai.onnx:Mod",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeMod(ctx, node, node.input(0).c_str(),
                               node.input(1).c_str());
       }},
      {"ai.onnx:Mish",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeMish(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Multinomial",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         generator::ComputeShapeMultinomial(ctx, node);
       }},
      {"ai.onnx:Neg",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeNeg(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Not",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         logical::ComputeShapeNot(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:PRelu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapePRelu(ctx, node, node.input(0).c_str(),
                                 node.input(1).c_str());
       }},
      {"ai.onnx:Pad",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapePad(ctx, node);
       }},
      {"ai.onnx:Pow",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapePow(ctx, node, node.input(0).c_str(),
                               node.input(1).c_str());
       }},
      {"ai.onnx:Sub",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeSub(ctx, node, node.input(0).c_str(),
                               node.input(1).c_str());
       }},
      {"ai.onnx:MatMul",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeMatMul(ctx, node, node.input(0).c_str(),
                                  node.input(1).c_str());
       }},
      {"ai.onnx:MatMulInteger",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         math::ComputeShapeMatMulInteger(ctx, node, node.input(0).c_str(),
                                         node.input(1).c_str());
       }},
      {"ai.onnx:Max",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeMax(ctx, node);
       }},
      {"ai.onnx:Mean",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeMean(ctx, node);
       }},
      {"ai.onnx:MelWeightMatrix",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 5);
         generator::ComputeShapeMelWeightMatrix(ctx, node);
       }},
      {"ai.onnx:Min",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeMin(ctx, node);
       }},
      {"ai.onnx:Optional",
       [](ShapesContext &ctx, const NodeProto &node) {
         optional::ComputeShapeOptional(ctx, node);
       }},
      {"ai.onnx:OptionalGetElement",
       [](ShapesContext &ctx, const NodeProto &node) {
         optional::ComputeShapeOptionalGetElement(ctx, node);
       }},
      {"ai.onnx:OptionalHasElement",
       [](ShapesContext &ctx, const NodeProto &node) {
         optional::ComputeShapeOptionalHasElement(ctx, node);
       }},
      {"ai.onnx:QuantizeLinear",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         const std::string x_name = node.input(0);
         const std::string zp_name =
             node.input_size() >= 3 ? std::string(node.input(2)) : std::string();
         quantization::ComputeShapeQuantizeLinear(ctx, node, x_name.c_str(),
                                                  zp_name.empty() ? nullptr : zp_name.c_str());
       }},
      {"ai.onnx:DynamicQuantizeLinear",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& x_name = node.input(0);
         quantization::ComputeShapeDynamicQuantizeLinear(ctx, node, x_name.c_str());
       }},
      {"ai.onnx:QLinearConv",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 8);
         const std::string& x_name = node.input(0);
         const std::string& w_name = node.input(3);
         const std::string& yzp_name = node.input(7);
         quantization::ComputeShapeQLinearConv(ctx, node, x_name.c_str(), w_name.c_str(),
                                               yzp_name.c_str());
       }},
      {"ai.onnx:QLinearMatMul",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 8);
         const std::string& a_name = node.input(0);
         const std::string& b_name = node.input(3);
         const std::string& yzp_name = node.input(7);
         quantization::ComputeShapeQLinearMatMul(ctx, node, a_name.c_str(), b_name.c_str(),
                                                 yzp_name.c_str());
       }},
      {"ai.onnx:ReduceSum",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceSum(ctx, node, data_name.c_str(),
                                          node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceSumSquare",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceSumSquare(
             ctx, node, data_name.c_str(),
             node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceL1",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceL1(ctx, node, data_name.c_str(),
                                         node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceL2",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceL2(ctx, node, data_name.c_str(),
                                         node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceLogSum",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceLogSum(
             ctx, node, data_name.c_str(),
             node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceLogSumExp",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceLogSumExp(
             ctx, node, data_name.c_str(),
             node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceMax",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceMax(ctx, node, data_name.c_str(),
                                          node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceMean",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceMean(ctx, node, data_name.c_str(),
                                           node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceMin",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceMin(ctx, node, data_name.c_str(),
                                          node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:ReduceProd",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         const std::string& data_name = node.input(0);
         const std::string& axes_name =
             node.input_size() >= 2 ? node.input(1) : utils::String::empty_string();
         reduction::ComputeShapeReduceProd(ctx, node, data_name.c_str(),
                                           node.input_size() >= 2 ? axes_name.c_str() : nullptr);
       }},
      {"ai.onnx:Reshape",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeReshape(ctx, node);
       }},
      {"ai.onnx:DepthToSpace",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeDepthToSpace(ctx, node);
       }},
      {"ai.onnx:SpaceToDepth",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeSpaceToDepth(ctx, node);
       }},
      {"ai.onnx:Slice",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         tensor::ComputeShapeSlice(ctx, node);
       }},
      {"ai.onnx:Split",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeSplit(ctx, node);
       }},
      {"ai.onnx:Tile",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeTile(ctx, node);
       }},
      {"ai.onnx:Squeeze",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeSqueeze(ctx, node);
       }},
      {"ai.onnx:Transpose",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeTranspose(ctx, node);
       }},
      {"ai.onnx:Trilu",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeTrilu(ctx, node);
       }},
      {"ai.onnx:CenterCropPad",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeCenterCropPad(ctx, node);
       }},
      {"ai.onnx:ReverseSequence",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeReverseSequence(ctx, node);
       }},
      {"ai.onnx:Unsqueeze",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         tensor::ComputeShapeUnsqueeze(ctx, node);
       }},
      {"ai.onnx:Upsample",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeUpsample(ctx, node);
       }},
      {"ai.onnx:Resize",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeResize(ctx, node);
       }},
      {"ai.onnx:NonZero",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeNonZero(ctx, node);
       }},
      {"ai.onnx:OneHot",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         tensor::ComputeShapeOneHot(ctx, node);
       }},
      {"ai.onnx:Unique",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeUnique(ctx, node);
       }},
      {"ai.onnx:Shape",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeShape(ctx, node);
       }},
      {"ai.onnx:Size",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeSize(ctx, node);
       }},
      {"ai.onnx:Identity",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         tensor::ComputeShapeIdentity(ctx, node);
       }},
      {"ai.onnx:RoiAlign",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 3);
        nn::ComputeShapeRoiAlign(ctx, node, node.input(0).c_str(),
                                 node.input(1).c_str(),
                                  node.input(2).c_str());
       }},
      {"ai.onnx:RotaryEmbedding",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         nn::ComputeShapeRotaryEmbedding(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:CausalConvWithState",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         nn::ComputeShapeCausalConvWithState(ctx, node, node.input(0).c_str(),
                                             node.input(1).c_str());
       }},
      {"ai.onnx:NonMaxSuppression",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         nn::ComputeShapeNonMaxSuppression(ctx, node, node.input(0).c_str(),
                                           node.input(1).c_str());
       }},
      {"ai.onnx:RNN",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         const std::string x_name = node.input(0);
         const std::string r_name = node.input(2);
         nn::ComputeShapeRNN(ctx, node, x_name.c_str(), r_name.empty() ? nullptr : r_name.c_str());
       }},
      {"ai.onnx:GRU",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         const std::string x_name = node.input(0);
         const std::string r_name = node.input(2);
         nn::ComputeShapeRNN(ctx, node, x_name.c_str(), r_name.empty() ? nullptr : r_name.c_str());
       }},
      {"ai.onnx:LSTM",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 3);
         const std::string x_name = node.input(0);
         const std::string r_name = node.input(2);
         nn::ComputeShapeRNN(ctx, node, x_name.c_str(), r_name.empty() ? nullptr : r_name.c_str());
       }},
      {"ai.onnx:SequenceConstruct",
       [](ShapesContext &ctx, const NodeProto &node) {
         sequence::ComputeShapeSequenceConstruct(ctx, node);
       }},
      {"ai.onnx:SequenceEmpty",
       [](ShapesContext &ctx, const NodeProto &node) {
         sequence::ComputeShapeSequenceEmpty(ctx, node);
       }},
      {"ai.onnx:SequenceLength",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         sequence::ComputeShapeSequenceLength(ctx, node);
       }},
      {"ai.onnx:SequenceErase",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         sequence::ComputeShapeSequenceErase(ctx, node);
       }},
      {"ai.onnx:SequenceAt",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         sequence::ComputeShapeSequenceAt(ctx, node);
       }},
      {"ai.onnx:SequenceInsert",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         sequence::ComputeShapeSequenceInsert(ctx, node);
       }},
      {"ai.onnx:SequenceMap",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         sequence::ComputeShapeSequenceMap(ctx, node);
       }},
      {"ai.onnx:SplitToSequence",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         sequence::ComputeShapeSplitToSequence(ctx, node);
       }},
      {"ai.onnx:ConcatFromSequence",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         sequence::ComputeShapeConcatFromSequence(ctx, node);
       }},
      {"ai.onnx:Sigmoid",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSigmoid(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Softmax",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSoftmax(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:LogSoftmax",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeLogSoftmax(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Softplus",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSoftplus(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Softsign",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeSoftsign(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:Shrink",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         math::ComputeShapeShrink(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:SoftmaxCrossEntropyLoss",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         const std::string& scores_name = node.input(0);
         const std::string& labels_name = node.input(1);
         const std::string& weights_name =
             node.input_size() >= 3 ? node.input(2) : utils::String::empty_string();
         math::ComputeShapeSoftmaxCrossEntropyLoss(
             ctx, node, scores_name.c_str(), labels_name.c_str(),
             weights_name.empty() ? nullptr : weights_name.c_str());
       }},
      {"ai.onnx:NegativeLogLikelihoodLoss",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         const std::string& input_name = node.input(0);
         const std::string& target_name = node.input(1);
         const std::string& weight_name =
             node.input_size() >= 3 ? node.input(2) : utils::String::empty_string();
         math::ComputeShapeNegativeLogLikelihoodLoss(
             ctx, node, input_name.c_str(), target_name.c_str(),
             weight_name.empty() ? nullptr : weight_name.c_str());
       }},
      {"ai.onnx:StringConcat",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         text::ComputeShapeStringConcat(ctx, node, node.input(0).c_str(),
                                        node.input(1).c_str());
       }},
      {"ai.onnx:StringSplit",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         text::ComputeShapeStringSplit(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:StringNormalizer",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         text::ComputeShapeStringNormalizer(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:RegexFullMatch",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         text::ComputeShapeRegexFullMatch(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx:TfIdfVectorizer",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         text::ComputeShapeTfIdfVectorizer(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:Binarizer",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeBinarizer(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:CastMap",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeCastMap(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:CategoryMapper",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeCategoryMapper(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:DictVectorizer",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeDictVectorizer(ctx, node,
                                                   node.input(0).c_str());
       }},
      {"ai.onnx.ml:FeatureVectorizer",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         std::vector<std::string> inputs;
         inputs.reserve(static_cast<size_t>(node.input_size()));
         for (int i = 0; i < node.input_size(); ++i) {
           inputs.emplace_back(node.input(i));
         }
         traditionalml::ComputeShapeFeatureVectorizer(ctx, node, inputs);
       }},
      {"ai.onnx.ml:Imputer",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeImputer(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:ArrayFeatureExtractor",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 2);
         traditionalml::ComputeShapeArrayFeatureExtractor(
             ctx, node, node.input(0).c_str(), node.input(1).c_str());
       }},
      {"ai.onnx.ml:LabelEncoder",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeLabelEncoder(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:LinearClassifier",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeLinearClassifier(ctx, node,
                                                    node.input(0).c_str());
       }},
      {"ai.onnx.ml:LinearRegressor",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeLinearRegressor(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:Normalizer",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeNormalizer(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:OneHotEncoder",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeOneHotEncoder(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:Scaler",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 1);
         traditionalml::ComputeShapeScaler(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:SVMClassifier",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        traditionalml::ComputeShapeSVMClassifier(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:SVMRegressor",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        traditionalml::ComputeShapeSVMRegressor(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:TreeEnsemble",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        traditionalml::ComputeShapeTreeEnsemble(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.ml:TreeEnsembleClassifier",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        traditionalml::ComputeShapeTreeEnsembleClassifier(ctx, node,
                                                         node.input(0).c_str());
       }},
      {"ai.onnx.ml:TreeEnsembleRegressor",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        traditionalml::ComputeShapeTreeEnsembleRegressor(ctx, node,
                                                        node.input(0).c_str());
       }},
      {"ai.onnx.ml:ZipMap",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 1);
        traditionalml::ComputeShapeZipMap(ctx, node, node.input(0).c_str());
       }},
      {"ai.onnx.preview:FlexAttention",
       [](ShapesContext &ctx, const NodeProto &node) {
        RequireInputs(node, 3);
        preview::ComputeShapeFlexAttention(ctx, node, node.input(0).c_str(),
                                           node.input(1).c_str(),
                                            node.input(2).c_str());
       }},
      {"ai.onnx.preview.training:Adagrad",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 5);
         training::ComputeShapeAdagrad(ctx, node);
       }},
      {"ai.onnx.preview.training:Adam",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 6);
         training::ComputeShapeAdam(ctx, node);
       }},
      {"ai.onnx.preview.training:Momentum",
       [](ShapesContext &ctx, const NodeProto &node) {
         RequireInputs(node, 5);
         training::ComputeShapeMomentum(ctx, node);
       }},
      {"ai.rt:DelayedInitializer",
       [](ShapesContext &ctx, const NodeProto &node) {
        rt::ComputeShapeDelayedInitializer(ctx, node);
       }},
  };
  return table;
}

using ::onnx_light::core::shapes::ComputePeakMemoryFn;

// Built-in table of every ``onnx_shapes`` peak-memory function, keyed by
// ``"<domain>:<op_type>"``. Mirrors :cpp:func:`BuiltinShapeFunctions`; only
// used to populate the shared ``core::shapes`` peak-memory dispatch table via
// :cpp:func:`RegisterPeakMemoryFunctions`. Operators without an entry report a
// peak memory of ``0`` through :cpp:func:`core::shapes::ComputePeakMemory`.
const std::unordered_map<std::string, ComputePeakMemoryFn> &BuiltinPeakMemoryFunctions() {
  static const std::unordered_map<std::string, ComputePeakMemoryFn> table = {
      {"ai.onnx:Attention", nn::ComputePeakMemoryAttention},
  };
  return table;
}

} // namespace
} // namespace shapes

void RegisterShapeFunctions() {
  // `lib_onnx_shape` is a plain static archive: a translation unit with no
  // symbol referenced from elsewhere would simply be dropped by the linker,
  // so registration cannot rely on a file-scope static object running as a
  // side effect of "just being linked in" (unlike, say,
  // `OpSchemaRegistry::map()` in `onnx_lib`, which can call its registration
  // functions lazily because both live in the same library). Callers must
  // invoke this function explicitly before using the shape-inference engine;
  // it is idempotent, so calling it more than once (or from multiple
  // independent entry points) is safe and cheap after the first call.
  static const bool kRegistered = [] {
    for (const auto &entry : shapes::BuiltinShapeFunctions()) {
      const std::string &key = entry.first;
      const std::size_t sep = key.find(':');
      ::onnx_light::core::shapes::RegisterComputeShapeFn(key.substr(0, sep), key.substr(sep + 1),
                                                         entry.second);
    }
    return true;
  }();
  (void)kRegistered;
}

void RegisterPeakMemoryFunctions() {
  // Mirrors :cpp:func:`RegisterShapeFunctions` for the peak-memory dispatch
  // table. Idempotent: the registration work runs once, guarded by a
  // function-local static, so repeated or multi-entry-point calls are safe.
  static const bool kRegistered = [] {
    for (const auto &entry : shapes::BuiltinPeakMemoryFunctions()) {
      const std::string &key = entry.first;
      const std::size_t sep = key.find(':');
      ::onnx_light::core::shapes::RegisterComputePeakMemoryFn(
          key.substr(0, sep), key.substr(sep + 1), ::onnx_light::core::shapes::Device::kCPU,
          entry.second);
    }
    return true;
  }();
  (void)kRegistered;
}

} // namespace ONNX_LIGHT_NAMESPACE::onnx_shapes
