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

#pragma once

#include "onnx_core/shapes/dispatch_table.h"
#include "onnx_core/shapes/shape_broadcast.h"
#include "onnx_core/shapes/shape_check.h"
#include "onnx_core/shapes/shape_inference.h"
#include "onnx_core/shapes/shapes_context.h"
#include "onnx_core/symbolic/sym_map.h"
#include "onnx_core/symbolic/sym_sequence.h"
#include "onnx_core/symbolic/sym_tensor.h"
#include "onnx_proto/onnx.h"

#include <cstdint>
#include <vector>

/**
 * @file shape_nn.h
 * @brief Shape-inference functions for ONNX operators in the ``nn`` family.
 */

namespace ONNX_LIGHT_NAMESPACE::onnx_shapes::shapes {

// The generic shape-inference engine (ShapesContext, dispatch table,
// domain constants, ...) lives in ``onnx_core`` so it never depends on
// any particular set of operator implementations; the symbolic value
// descriptors (SymDim, SymShape, SymTensor, ...) live in
// ``onnx_core::symbolic`` for the same reason. Bring them into
// ``onnx_shapes::shapes`` so shape functions can keep referring to them
// unqualified, exactly as before those types moved to ``onnx_core``.
using ::onnx_light::core::shapes::BroadcastDimOp;
using ::onnx_light::core::shapes::BroadcastShapes;
using ::onnx_light::core::shapes::CheckNodeOpAndOutput;
using ::onnx_light::core::shapes::ComputeShapeBinaryBroadcast;
using ::onnx_light::core::shapes::InferShapesModel;
using ::onnx_light::core::shapes::kOnnxDomain;
using ::onnx_light::core::shapes::kUnknownOpsetVersion;
using ::onnx_light::core::shapes::PropagateValueAsShapeArithmetic;
using ::onnx_light::core::shapes::ShapeEvent;
using ::onnx_light::core::shapes::ShapeEventAction;
using ::onnx_light::core::shapes::ShapeEventActionName;
using ::onnx_light::core::shapes::ShapesContext;

using ::onnx_light::core::symbolic::DataTypeToTensorType;
using ::onnx_light::core::symbolic::Device;
using ::onnx_light::core::symbolic::GPUIndex;
using ::onnx_light::core::symbolic::IsGPU;
using ::onnx_light::core::symbolic::IsIntegerTensorType;
using ::onnx_light::core::symbolic::kMaxGPUIndex;
using ::onnx_light::core::symbolic::kMaxOptimRank;
using ::onnx_light::core::symbolic::kOptimValueAsShapeMaxElements;
using ::onnx_light::core::symbolic::kValueInfoDeviceMetadataKey;
using ::onnx_light::core::symbolic::kValueInfoMaxMetadataKey;
using ::onnx_light::core::symbolic::kValueInfoMinMetadataKey;
using ::onnx_light::core::symbolic::ShapeFromTensorProtoDims;
using ::onnx_light::core::symbolic::SymCmpResult;
using ::onnx_light::core::symbolic::SymDim;
using ::onnx_light::core::symbolic::SymMap;
using ::onnx_light::core::symbolic::SymSequence;
using ::onnx_light::core::symbolic::SymShape;
using ::onnx_light::core::symbolic::SymTensor;
using ::onnx_light::core::symbolic::TensorType;
using ::onnx_light::core::symbolic::TensorTypeToDataType;

namespace nn {

/**
 * Computes the output :cpp:class:`SymTensor` of an ``AveragePool`` node
 * and stores it in ``ctx``.
 *
 * The output dtype matches the input dtype. The output rank is the input
 * rank (``[N, C, D1, ..., Dk]``). The first two output dimensions (``N``
 * and ``C``) are copied from the input. For each spatial axis ``i`` the
 * output dimension is computed from the ``kernel_shape``, ``strides``,
 * ``pads``, and ``ceil_mode`` attributes using the same rule as
 * :cpp:func:`onnx_kernels::kernel::AveragePool` and ONNX Runtime:
 * when ``ceil_mode=1`` and the last sliding window would start entirely
 * in the right padded region, it is dropped. ``auto_pad`` other than the
 * default ``"NOTSET"`` (or ``"VALID"``) is not supported, and symbolic
 * spatial dimensions propagate symbolically.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``;
 *              on return it also contains an entry for ``node.output(0)``.
 * @param node  The ``AveragePool`` ``NodeProto`` whose output should be
 *              described. ``node.op_type()`` must be ``"AveragePool"``
 *              and ``node`` must declare at least one output.
 * @param x     Name of the input value to read from ``ctx``. Must be
 *              present in ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"AveragePool"``, if ``node`` has no output, if the input
 *         rank is inconsistent with the (required) ``kernel_shape``
 *         attribute, or if ``auto_pad`` is set to a value other than
 *         ``"NOTSET"`` / ``"VALID"`` (only explicit ``pads`` are
 *         supported).
 * @throws std::out_of_range     if ``x`` is not present in ``ctx``.
 */
void ComputeShapeAveragePool(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of an ``LpPool`` node and
 * stores it in ``ctx``.
 *
 * The output dtype matches the input dtype and the output shape is computed
 * identically to :cpp:func:`ComputeShapeAveragePool` (the two operators share
 * the same ``kernel_shape``/``strides``/``pads``/``dilations``/``ceil_mode``
 * / ``auto_pad`` attribute schema; ``LpPool`` has an extra ``p`` attribute
 * that does not affect the output shape).
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``;
 *              on return it also contains an entry for ``node.output(0)``.
 * @param node  The ``LpPool`` ``NodeProto`` whose output should be
 *              described. ``node.op_type()`` must be ``"LpPool"`` and
 *              ``node`` must declare at least one output.
 * @param x     Name of the input value to read from ``ctx``. Must be
 *              present in ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not ``"LpPool"``,
 *         if ``node`` has no output, if the input rank is inconsistent
 *         with the (required) ``kernel_shape`` attribute, or if
 *         ``auto_pad`` is set to an unsupported value.
 * @throws std::out_of_range     if ``x`` is not present in ``ctx``.
 */
void ComputeShapeLpPool(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of a global pooling node
 * (``GlobalAveragePool``, ``GlobalMaxPool``, or ``GlobalLpPool``) and stores
 * it in ``ctx``.
 *
 * The output dtype matches the input dtype. The output rank equals the input
 * rank. The first two output dimensions (``N`` and ``C``) are copied from the
 * input; all remaining spatial dimensions are set to ``1``.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``;
 *              on return it also contains an entry for ``node.output(0)``.
 * @param node  The global pooling ``NodeProto`` whose output should be
 *              described. ``node.op_type()`` must be one of
 *              ``"GlobalAveragePool"``, ``"GlobalMaxPool"`` or
 *              ``"GlobalLpPool"`` and ``node`` must declare at least one
 *              output.
 * @param x     Name of the input value to read from ``ctx``. Must be
 *              present in ``ctx``.
 *
 * @throws std::invalid_argument if the input rank is less than 2 or if
 *         ``node`` has no output.
 * @throws std::out_of_range     if ``x`` is not present in ``ctx``.
 */
void ComputeShapeGlobalPool(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``Flatten`` node and
 * stores it in ``ctx``.
 *
 * The output dtype matches the input dtype. The output shape is always
 * rank 2: ``(prod(input.shape[0:axis]), prod(input.shape[axis:rank]))``,
 * where ``axis`` is the integer attribute (default ``1``) and may be
 * negative (counted from the back). Symbolic dimensions propagate
 * symbolically — when any contributing dim is symbolic the corresponding
 * output dim becomes a fresh symbolic expression.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``;
 *              on return it also contains an entry for ``node.output(0)``.
 * @param node  The ``Flatten`` ``NodeProto`` whose output should be
 *              described. ``node.op_type()`` must be ``"Flatten"`` and
 *              ``node`` must declare at least one output.
 * @param x     Name of the input value to read from ``ctx``. Must be
 *              present in ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not ``"Flatten"``,
 *         if ``node`` has no output, or if ``axis`` is out of range
 *         ``[-r, r]`` for input rank ``r``.
 * @throws std::out_of_range     if ``x`` is not present in ``ctx``.
 */
void ComputeShapeFlatten(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of an ``LRN`` (Local Response
 * Normalization) node and stores it in ``ctx``.
 *
 * The output dtype and shape match the input dtype and shape. The input must
 * have rank ``>= 2`` (``N``, ``C``, ``D1``, ..., ``Dk``).
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``;
 *              on return it also contains an entry for ``node.output(0)``.
 * @param node  The ``LRN`` ``NodeProto`` whose output should be described.
 *              ``node.op_type()`` must be ``"LRN"`` and ``node`` must declare
 *              at least one output.
 * @param x     Name of the input value to read from ``ctx``. Must be
 *              present in ``ctx``.
 *
 * @throws std::invalid_argument if the input rank is less than 2 or if
 *         ``node`` has no output.
 * @throws std::out_of_range     if ``x`` is not present in ``ctx``.
 */
void ComputeShapeLRN(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of an ``LpNormalization`` node
 * and stores it in ``ctx``.
 *
 * The output dtype and shape match the input dtype and shape. The input must
 * have rank ``>= 1``.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``;
 *              on return it also contains an entry for ``node.output(0)``.
 * @param node  The ``LpNormalization`` ``NodeProto`` whose output should be
 *              described. ``node.op_type()`` must be ``"LpNormalization"`` and
 *              ``node`` must declare at least one output.
 * @param x     Name of the input value to read from ``ctx``. Must be
 *              present in ``ctx``.
 *
 * @throws std::invalid_argument if the input has rank 0 or if ``node`` has
 *         no output.
 * @throws std::out_of_range     if ``x`` is not present in ``ctx``.
 */
void ComputeShapeLpNormalization(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``BatchNormalization``
 * node and stores it in ``ctx``.
 *
 * The op has between 1 and 5 outputs depending on opset / mode. The first
 * output ``Y`` always has the same dtype and shape as ``X``. The additional
 * outputs are 1-D tensors of size equal to the channel dimension ``C`` (the
 * dim at index 1 of ``X`` when ``X`` has rank >= 2; ``1`` otherwise, per the
 * opset-9+ rule). Output dtypes follow the upstream spec:
 *
 * - opset 1..9: every extra output uses the same dtype as ``X`` (type ``T``).
 * - opset 14: outputs 1/2 (``running_mean`` / ``running_var``) take the dtype
 *   of ``input_mean`` / ``input_var`` (type ``U``).
 * - opset 15+: outputs 1/2 (``running_mean`` / ``running_var``) take the
 *   dtype of ``input_mean`` / ``input_var`` (type ``T2``).
 *
 * The ``training_mode`` attribute (opset 14+) and the legacy ``is_test``
 * attribute (opset 1, 6) are honored to decide whether extra outputs are
 * required. Only the channel-dim of the secondary outputs is inferred; for
 * a symbolic channel dim a fresh symbolic expression
 * ``BatchNormalization.C(<expr>)`` is propagated.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``;
 *              on return it also contains an entry for each declared output
 *              of ``node``.
 * @param node  The ``BatchNormalization`` ``NodeProto``. ``node.op_type()``
 *              must be ``"BatchNormalization"`` and at least one output must
 *              be declared.
 * @param x     Name of the data input value to read from ``ctx``.
 * @param input_mean Name of the mean input value (used to obtain the dtype
 *                   of the secondary outputs from opset 14 onward). May be
 *                   ``nullptr`` if not needed (opset < 14 fallback).
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"BatchNormalization"`` or if ``node`` has no output.
 * @throws std::out_of_range     if ``x`` or ``input_mean`` is not present
 *                               in ``ctx`` when required.
 */
void ComputeShapeBatchNormalization(ShapesContext &ctx, const NodeProto &node, const char *x,
                                    const char *input_mean);

/**
 * Computes the output :cpp:class:`SymTensor` of an ``InstanceNormalization``
 * node and stores it in ``ctx``.
 *
 * The output dtype and shape are always inherited from input ``X``. ``scale``
 * and ``B`` inputs are not read for shape inference.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``; on
 *              return it also contains an entry for ``node.output(0)``.
 * @param node  The ``InstanceNormalization`` ``NodeProto``. ``node.op_type()``
 *              must be ``"InstanceNormalization"`` and ``node`` must declare
 *              at least one output.
 * @param x     Name of the data input value to read from ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"InstanceNormalization"`` or if ``node`` has no output.
 * @throws std::out_of_range if ``x`` is not present in ``ctx``.
 */
void ComputeShapeInstanceNormalization(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``GroupNormalization``
 * node and stores it in ``ctx``.
 *
 * The output dtype and shape are always inherited from input ``X``. ``scale``
 * and ``bias`` inputs are not read for shape inference.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``; on
 *              return it also contains an entry for ``node.output(0)``.
 * @param node  The ``GroupNormalization`` ``NodeProto``. ``node.op_type()``
 *              must be ``"GroupNormalization"`` and ``node`` must declare at
 *              least one output.
 * @param x     Name of the data input value to read from ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"GroupNormalization"`` or if ``node`` has no output.
 * @throws std::out_of_range if ``x`` is not present in ``ctx``.
 */
void ComputeShapeGroupNormalization(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of a
 * ``MeanVarianceNormalization`` node and stores it in ``ctx``.
 *
 * The output dtype and shape are always inherited from input ``X``.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``; on
 *              return it also contains an entry for ``node.output(0)``.
 * @param node  The ``MeanVarianceNormalization`` ``NodeProto``.
 * @param x     Name of the data input value to read from ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"MeanVarianceNormalization"`` or if ``node`` has no output.
 * @throws std::out_of_range if ``x`` is not present in ``ctx``.
 */
void ComputeShapeMeanVarianceNormalization(ShapesContext &ctx, const NodeProto &node,
                                           const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of an ``RMSNormalization``
 * node and stores it in ``ctx``.
 *
 * The output dtype and shape are always inherited from input ``X``. The
 * ``scale`` input is not read for shape inference.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``; on
 *              return it also contains an entry for ``node.output(0)``.
 * @param node  The ``RMSNormalization`` ``NodeProto``. ``node.op_type()``
 *              must be ``"RMSNormalization"`` and ``node`` must declare at
 *              least one output.
 * @param x     Name of the data input value to read from ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"RMSNormalization"`` or if ``node`` has no output.
 * @throws std::out_of_range if ``x`` is not present in ``ctx``.
 */
void ComputeShapeRMSNormalization(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``LayerNormalization``
 * node and stores it in ``ctx``.
 *
 * Output 0 (``Y``) always inherits ``X``'s dtype and shape. The optional
 * outputs ``Mean`` (output 1) and ``InvStdDev`` (output 2) have dtype
 * ``stash_type`` (default ``FLOAT``) and shape
 * ``[d[0], ..., d[axis-1], 1, ..., 1]`` (rank == rank(X)).
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``; on
 *              return it also contains an entry for every non-empty
 *              ``node.output(i)`` (``i = 0, 1, 2``).
 * @param node  The ``LayerNormalization`` ``NodeProto``. ``node.op_type()``
 *              must be ``"LayerNormalization"`` and ``node`` must declare
 *              at least one output.
 * @param x     Name of the data input value to read from ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"LayerNormalization"`` or if ``node`` has no output.
 * @throws std::out_of_range if ``x`` is not present in ``ctx``.
 */
void ComputeShapeLayerNormalization(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor`(s) of a ``Dropout`` node and
 * stores them in ``ctx``.
 *
 * ``output`` always has the same dtype and shape as ``data``. If the optional
 * second output ``mask`` is present and non-empty, it has dtype
 * ``TensorType::kBool`` and the same shape as ``data``. Optional inputs
 * ``ratio`` and ``training_mode`` must be scalars when present.
 */
void ComputeShapeDropout(ShapesContext &ctx, const NodeProto &node, const char *data,
                         const char *ratio = nullptr, const char *training_mode = nullptr);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``RoiAlign`` node and
 * stores it in ``ctx``.
 *
 * The output dtype matches the input feature-map dtype. The output shape
 * is ``(num_rois, C, output_height, output_width)`` where ``num_rois`` is
 * taken from dim 0 of ``rois`` (or, when that dim is symbolic and
 * ``batch_indices`` dim 0 is static, from ``batch_indices``); ``C`` is
 * taken from dim 1 of ``x``; and the spatial sizes come from the
 * ``output_height`` / ``output_width`` attributes (default 1).
 *
 * @param ctx              In/out context. Must already contain entries for
 *                         ``x``, ``rois``, and ``batch_indices``; on return
 *                         it also contains an entry for ``node.output(0)``.
 * @param node             The ``RoiAlign`` ``NodeProto`` whose output should
 *                         be described. ``node.op_type()`` must be
 *                         ``"RoiAlign"`` and ``node`` must declare at least
 *                         one output.
 * @param x                Name of the feature-map input value (rank 4) to
 *                         read from ``ctx``. Must be present in ``ctx``.
 * @param rois             Name of the RoIs input value (rank 2) to read
 *                         from ``ctx``. Must be present in ``ctx``.
 * @param batch_indices    Name of the batch-indices input value (rank 1) to
 *                         read from ``ctx``. Must be present in ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not ``"RoiAlign"``,
 *         if ``node`` has no output, if any input has the wrong rank, or
 *         if ``output_height`` / ``output_width`` is non-positive.
 * @throws std::out_of_range     if any input name is not present in ``ctx``.
 */
void ComputeShapeRoiAlign(ShapesContext &ctx, const NodeProto &node, const char *x,
                          const char *rois, const char *batch_indices);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``NonMaxSuppression`` node
 * and stores it in ``ctx``.
 *
 * The output dtype is always ``TensorType::kInt64``. The output shape is
 * always rank 2: ``(num_selected_indices, 3)`` where ``num_selected_indices``
 * is a fresh symbolic expression because it depends on the runtime values of
 * the inputs. The optional ``max_output_boxes_per_class``, ``iou_threshold``
 * and ``score_threshold`` inputs are not read by shape inference and may be
 * omitted from ``ctx``.
 *
 * @param ctx    In/out context. Must already contain entries for ``boxes`` and
 *               ``scores``; on return it also contains an entry for
 *               ``node.output(0)``.
 * @param node   The ``NonMaxSuppression`` ``NodeProto`` whose output should be
 *               described. ``node.op_type()`` must be ``"NonMaxSuppression"``
 *               and ``node`` must declare at least one output.
 * @param boxes  Name of the boxes input (rank 3, last dim == 4) in ``ctx``.
 * @param scores Name of the scores input (rank 3) in ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"NonMaxSuppression"``, if ``node`` has no output, or if the
 *         present inputs have inconsistent ranks/shapes.
 * @throws std::out_of_range     if ``boxes`` or ``scores`` is not present in
 *                               ``ctx``.
 */
void ComputeShapeNonMaxSuppression(ShapesContext &ctx, const NodeProto &node, const char *boxes,
                                   const char *scores);

/**
 * Computes the output :cpp:class:`SymTensor`(s) of an ``RNN``, ``GRU`` or
 * ``LSTM`` node and stores them in ``ctx``. The three operators share the
 * same output-shape semantics — only the number of outputs differs
 * (``RNN`` / ``GRU`` expose ``Y`` and ``Y_h``; ``LSTM`` also exposes
 * ``Y_c``).
 *
 * The output dtypes are inherited from ``X``. Shapes follow the upstream
 * ``RNNShapeInference`` rules:
 *
 * - Read ``direction`` (``"forward"`` (default), ``"reverse"`` or
 *   ``"bidirectional"``) to derive ``num_directions`` (1 or 2). Unknown
 *   values leave ``num_directions`` symbolic.
 * - Read ``hidden_size``; when missing or non-positive, fall back to
 *   ``R.shape[2]`` (the recurrence weight's last dim) when known.
 * - Read ``layout`` (default 0). For ``layout=0`` derive ``seq_length`` and
 *   ``batch_size`` from ``X.shape[0]`` and ``X.shape[1]``; for ``layout=1``
 *   the order is reversed.
 *
 * ``Y`` has rank 4 (``[seq_length, num_directions, batch_size,
 * hidden_size]`` for ``layout=0``; ``[batch_size, seq_length,
 * num_directions, hidden_size]`` for ``layout=1``). ``Y_h`` (and ``Y_c``,
 * for ``LSTM``) has rank 3 (``[num_directions, batch_size, hidden_size]``
 * for ``layout=0``; ``[batch_size, num_directions, hidden_size]`` for
 * ``layout=1``). Missing output dims propagate as a fresh symbolic
 * expression labeled with the operator and field name.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``;
 *              ``r`` is consulted only as a fallback for ``hidden_size``
 *              and may be ``nullptr`` or missing from ``ctx``. On return
 *              ``ctx`` also contains an entry for each declared (non-empty)
 *              output of ``node``.
 * @param node  The ``RNN`` / ``GRU`` / ``LSTM`` ``NodeProto``.
 *              ``node.op_type()`` must be one of ``"RNN"``, ``"GRU"`` or
 *              ``"LSTM"`` and at least one output must be declared.
 * @param x     Name of the data input value to read from ``ctx``. Must be
 *              present in ``ctx`` and have rank 3.
 * @param r     Name of the recurrence-weight input value (used as a
 *              fallback source of ``hidden_size``). May be ``nullptr`` or
 *              absent from ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not one of
 *         ``"RNN"`` / ``"GRU"`` / ``"LSTM"``, if ``node`` has no output,
 *         or if ``X`` does not have rank 3.
 * @throws std::out_of_range     if ``x`` is not present in ``ctx``.
 */
void ComputeShapeRNN(ShapesContext &ctx, const NodeProto &node, const char *x, const char *r);

/**
 * Computes the output :cpp:class:`SymTensor`(s) of an ``Attention`` node
 * (since opset 23 in the ``ai.onnx`` domain) and stores them in ``ctx``.
 *
 * ``Attention`` accepts 3 to 7 inputs (``Q``, ``K``, ``V`` and the optional
 * ``attn_mask``, ``past_key``, ``past_value``, ``nonpad_kv_seqlen``) and
 * exposes between 1 and 4 outputs (``Y`` plus the optional ``present_key``,
 * ``present_value`` and ``qk_matmul_output``). Only the rank-4 input form
 * is described here; the rank-3 form (where ``q_num_heads`` /
 * ``kv_num_heads`` are read from attributes) defers shape inference to
 * the dispatcher.
 *
 * For rank-4 inputs the function infers:
 *
 *   - ``Y``: ``(batch_size, q_num_heads, q_sequence_length, v_head_size)``
 *     with dtype matching ``Q``.
 *   - ``present_key``: ``(batch_size, kv_num_heads, total_sequence_length,
 *     head_size)`` where ``total_sequence_length = past_sequence_length +
 *     kv_sequence_length`` (or just ``kv_sequence_length`` when no past
 *     state is provided).
 *   - ``present_value``: ``(batch_size, kv_num_heads,
 *     total_sequence_length, v_head_size)`` with dtype matching ``V``.
 *   - ``qk_matmul_output``: ``(batch_size, q_num_heads, q_sequence_length,
 *     total_sequence_length)`` with dtype matching ``Q``.
 *
 * ``q_num_heads`` must be a multiple of ``kv_num_heads`` when both are
 * static (Grouped Query Attention). Symbolic dimensions propagate
 * symbolically.
 *
 * @param ctx          In/out context. Must already contain entries for
 *                     ``q``, ``k`` and ``v``; on return it also contains
 *                     an entry for each declared output of ``node``.
 * @param node         The ``Attention`` ``NodeProto`` whose outputs should
 *                     be described. ``node.op_type()`` must be
 *                     ``"Attention"`` and ``node`` must declare at least
 *                     one output.
 * @param q            Name of the query input value (rank 4).
 * @param k            Name of the key input value (rank 4).
 * @param v            Name of the value input value (rank 4).
 * @param past_key     Optional name of the past_key input (rank 4). When
 *                     not ``nullptr`` and present in ``ctx``, contributes
 *                     ``past_sequence_length`` to the present_* outputs.
 * @param past_value   Optional name of the past_value input (rank 4).
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"Attention"``, if ``node`` has no output, if any of ``Q``/
 *         ``K``/``V`` has rank other than 4, or if static shapes are
 *         inconsistent (mismatched batch size, mismatched head dim,
 *         mismatched kv_sequence_length, or ``q_num_heads`` not a
 *         multiple of ``kv_num_heads``).
 * @throws std::out_of_range     if ``q``/``k``/``v`` is not present in
 *                               ``ctx``.
 */
void ComputeShapeAttention(ShapesContext &ctx, const NodeProto &node, const char *q, const char *k,
                           const char *v, const char *past_key = nullptr,
                           const char *past_value = nullptr);

/**
 * Estimates the peak scratch memory (in bytes) required to evaluate an
 * ``Attention`` node, following the
 * :cpp:type:`core::shapes::ComputePeakMemoryFn` convention: it takes the
 * :cpp:enum:`Device` the operator runs on followed by the shapes of the
 * node's inputs.
 *
 * The estimate is dominated by the ``QK^T`` attention-score buffer, whose
 * logical shape is ``(batch, q_num_heads, q_sequence_length,
 * total_sequence_length)``. The number of elements in that buffer is
 * multiplied by 4 bytes (the ``float32`` accumulation type used for the
 * scores). Only the rank-4 ``(batch, num_heads, sequence, head_size)`` form
 * of ``Q``/``K`` is handled; when fewer than three inputs are provided, when
 * either input is not rank 4, or when any contributing dimension is symbolic,
 * the function returns ``0`` (no concrete estimate available).
 *
 * @param device       Device on which the operator executes (currently does
 *                     not affect the estimate).
 * @param input_shapes Shapes of the node's inputs, in order; ``input_shapes[0]``
 *                     is ``Q`` and ``input_shapes[1]`` is ``K``.
 * @returns The estimated peak scratch memory in bytes, or ``0`` when no
 *          concrete estimate can be produced.
 */
int64_t ComputePeakMemoryAttention(Device device, const std::vector<SymShape> &input_shapes);

/**
 * Computes the output :cpp:class:`SymTensor`(s) of a ``LinearAttention``
 * node (since opset 27 in the ``ai.onnx`` domain) and stores them in ``ctx``.
 *
 * ``LinearAttention`` takes 3 to 6 inputs (``query``, ``key``, ``value`` and
 * the optional ``past_state``, ``decay``, ``beta``) in rank-3 packed format
 * ``(batch_size, sequence_length, num_heads * head_size)``, plus the
 * required attributes ``q_num_heads`` and ``kv_num_heads``. It produces one
 * or two outputs (``output`` and the optional ``present_state``).
 *
 * Inferred shapes are:
 *
 *   - ``output``: ``(batch_size, sequence_length, q_num_heads * d_v)``
 *     with dtype matching ``query``, where ``d_v = value.shape[-1] /
 *     kv_num_heads``.
 *   - ``present_state``: ``(batch_size, kv_num_heads, d_k, d_v)`` where
 *     ``d_k = key.shape[-1] / kv_num_heads``. The dtype matches
 *     ``past_state`` when provided, otherwise ``query``.
 *
 * ``q_num_heads`` must be a positive multiple of ``kv_num_heads`` (Grouped
 * Query Attention). Symbolic dimensions propagate symbolically.
 *
 * @param ctx         In/out context. Must already contain entries for
 *                    ``query``, ``key`` and ``value``; on return it also
 *                    contains an entry for each declared output of ``node``.
 * @param node        The ``LinearAttention`` ``NodeProto`` whose outputs
 *                    should be described. ``node.op_type()`` must be
 *                    ``"LinearAttention"`` and ``node`` must declare at
 *                    least one output and the ``q_num_heads`` /
 *                    ``kv_num_heads`` attributes.
 * @param query       Name of the query input value (rank 3).
 * @param key         Name of the key input value (rank 3).
 * @param value       Name of the value input value (rank 3).
 * @param past_state  Optional name of the past_state input (rank 4). When
 *                    not ``nullptr`` and present in ``ctx``, refines the
 *                    inferred ``d_k`` / ``d_v`` dimensions and sets the
 *                    dtype of ``present_state``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"LinearAttention"``, if ``node`` has no output, if
 *         ``q_num_heads`` / ``kv_num_heads`` is missing or invalid, if
 *         any of ``query``/``key``/``value`` has rank other than 3, or
 *         if static shapes are inconsistent.
 * @throws std::out_of_range     if ``query``/``key``/``value`` is not
 *                               present in ``ctx``.
 */
void ComputeShapeLinearAttention(ShapesContext &ctx, const NodeProto &node, const char *query,
                                 const char *key, const char *value,
                                 const char *past_state = nullptr);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``RotaryEmbedding`` node
 * (since opset 23 in the ``ai.onnx`` domain) and stores it in ``ctx``.
 *
 * The output has the same shape and dtype as the input ``X``. ``X`` may be
 * either rank-4 ``(batch_size, num_heads, sequence_length, head_size)`` or
 * rank-3 ``(batch_size, sequence_length, hidden_size)``; in the latter case
 * the ``num_heads`` attribute must be set.
 *
 * @param ctx   In/out context. Must already contain an entry for ``x``; on
 *              return it also contains an entry for ``node.output(0)``.
 * @param node  The ``RotaryEmbedding`` ``NodeProto`` whose output should be
 *              described. ``node.op_type()`` must be ``"RotaryEmbedding"``
 *              and ``node`` must declare at least one output.
 * @param x     Name of the input value to read from ``ctx``. Must be present
 *              in ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"RotaryEmbedding"``, if ``node`` has no output, or if ``X`` has
 *         a rank other than 3 or 4.
 * @throws std::out_of_range     if ``x`` is not present in ``ctx``.
 */
void ComputeShapeRotaryEmbedding(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor`s of a ``CausalConvWithState``
 * node (since opset 27 in the ``ai.onnx`` domain) and stores them in ``ctx``.
 *
 * The first output has the same shape and dtype as the input. The second
 * output ``present_state`` has shape ``(batch_size, channels, k - 1)`` where
 * ``k`` is the kernel size (``weight.shape[2]``).
 *
 * @param ctx     In/out context. Must already contain entries for ``input``
 *                and ``weight``; on return it also contains entries for
 *                ``node.output(0)`` and ``node.output(1)``.
 * @param node    The ``CausalConvWithState`` ``NodeProto``. ``node.op_type()``
 *                must be ``"CausalConvWithState"`` and ``node`` must declare
 *                two outputs.
 * @param input   Name of the input data value (rank-3) to read from ``ctx``.
 * @param weight  Name of the weight value (rank-3) to read from ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"CausalConvWithState"``, if ``node`` declares less than two
 *         outputs, or if the input/weight ranks are not 3.
 * @throws std::out_of_range     if ``input`` or ``weight`` is not present
 *                               in ``ctx``.
 */
void ComputeShapeCausalConvWithState(ShapesContext &ctx, const NodeProto &node, const char *input,
                                     const char *weight);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``DeformConv`` node and
 * stores it in ``ctx``.
 *
 * The output dtype matches the input ``X`` dtype. The output shape is
 * ``(N, oC, o1, ..., on)`` where ``N`` is ``X.shape[0]``, ``oC`` is
 * ``W.shape[0]``, and each spatial dim ``oi`` is computed from the input
 * spatial dim, the (effective) kernel shape, ``strides``, ``pads`` and
 * ``dilations`` attributes — matching the upstream
 * ``convPoolShapeInference`` rule shared with ``Conv``. When the
 * ``kernel_shape`` attribute is missing, the kernel shape is taken from
 * ``W.shape[2..]``. Symbolic spatial dimensions are propagated symbolically.
 *
 * @param ctx   In/out context. Must already contain entries for ``x`` and
 *              ``w``; on return it also contains an entry for
 *              ``node.output(0)``.
 * @param node  The ``DeformConv`` ``NodeProto`` whose output should be
 *              described. ``node.op_type()`` must be ``"DeformConv"`` and
 *              ``node`` must declare at least one output.
 * @param x     Name of the input data value (rank >= 3) to read from
 *              ``ctx``. Must be present in ``ctx``.
 * @param w     Name of the weight value (rank >= 3) to read from ``ctx``.
 *              Must be present in ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not
 *         ``"DeformConv"``, if ``node`` has no output, if ``X`` and ``W``
 *         have inconsistent ranks, or if attributes have wrong sizes.
 * @throws std::out_of_range     if ``x`` or ``w`` is not present in ``ctx``.
 */
void ComputeShapeDeformConv(ShapesContext &ctx, const NodeProto &node, const char *x,
                            const char *w);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``Col2Im`` node and stores
 * it in ``ctx``.
 *
 * The output dtype matches the input ``input`` dtype. The output shape is
 * ``(N, C, dim_i1, ..., dim_iN)`` where ``N`` is ``input.shape[0]``, ``C`` is
 * ``input.shape[1] / product(block_shape)`` (when the ``block_shape``
 * initializer is known), and the spatial dimensions are taken from the
 * ``image_shape`` initializer when known. When either initializer is missing
 * or the corresponding input shape is symbolic, the affected dimensions are
 * propagated symbolically.
 *
 * @param ctx   In/out context. Must already contain entries for ``input``,
 *              ``image_shape`` and ``block_shape``; on return it also contains
 *              an entry for ``node.output(0)``.
 * @param node  The ``Col2Im`` ``NodeProto`` whose output should be described.
 *              ``node.op_type()`` must be ``"Col2Im"`` and ``node`` must
 *              declare at least one output.
 * @param input         Name of the data input value (rank 3) in ``ctx``.
 * @param image_shape   Name of the ``image_shape`` 1-D ``tensor(int64)`` value
 *                      in ``ctx``.
 * @param block_shape   Name of the ``block_shape`` 1-D ``tensor(int64)`` value
 *                      in ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not ``"Col2Im"``,
 *         if ``node`` has no output, or if the inputs have inconsistent
 *         shapes.
 * @throws std::out_of_range     if ``input``, ``image_shape`` or
 *                               ``block_shape`` is not present in ``ctx``.
 */
void ComputeShapeCol2Im(ShapesContext &ctx, const NodeProto &node, const char *input,
                        const char *image_shape, const char *block_shape);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``Conv`` node and stores
 * it in ``ctx``.
 *
 * The output dtype matches ``X``. The output shape is ``(N, M, o1, ..., on)``
 * where ``N`` is ``X.shape[0]``, ``M`` is ``W.shape[0]``, and each spatial
 * dim is derived from ``kernel_shape``, ``strides``, ``pads``, ``dilations``
 * and ``auto_pad`` following the upstream ``convPoolShapeInference`` rule.
 */
void ComputeShapeConv(ShapesContext &ctx, const NodeProto &node, const char *x, const char *w);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``ConvInteger`` node and
 * stores it in ``ctx``. The shape rule matches :cpp:func:`ComputeShapeConv`;
 * the output dtype is always ``TensorType::kInt32``.
 */
void ComputeShapeConvInteger(ShapesContext &ctx, const NodeProto &node, const char *x,
                             const char *w);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``ConvTranspose`` node
 * and stores it in ``ctx``.
 *
 * The output dtype matches ``X``. The output shape is ``(N, M, o1, ..., on)``
 * where ``M`` is ``W.shape[1] * group``. Each spatial dim is derived from
 * ``kernel_shape``, ``strides``, ``pads``, ``dilations``, ``output_padding``,
 * ``output_shape`` and ``auto_pad`` following the upstream
 * ``convTransposeShapeInference`` rule.
 */
void ComputeShapeConvTranspose(ShapesContext &ctx, const NodeProto &node, const char *x,
                               const char *w);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``MaxPool`` node.
 *
 * Mirrors :cpp:func:`ComputeShapeAveragePool`. The first output ``Y`` has
 * the same dtype as ``X`` and a shape derived from ``kernel_shape``,
 * ``strides``, ``pads``, ``dilations``, ``ceil_mode`` and ``auto_pad``.
 * If the node declares a second output (``Indices``) and that output's
 * name is non-empty, an entry is also added to ``ctx`` with dtype
 * ``int64`` and the same shape as ``Y``.
 */
void ComputeShapeMaxPool(ShapesContext &ctx, const NodeProto &node, const char *x);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``MaxUnpool`` node.
 *
 * The output dtype matches ``X``. When ``output_shape`` is provided and is
 * a rank-1 ``int64`` tensor with a known value, the output shape is taken
 * directly from those values; otherwise the output shape is computed from
 * ``kernel_shape``, ``strides`` and ``pads`` according to the formula
 * ``out_dim[i] = strides[i] * (in_dim[i] - 1) + kernel_shape[i] - pads[i] -
 * pads[i + k]``. ``output_shape`` may be ``nullptr`` for the two-input form.
 */
void ComputeShapeMaxUnpool(ShapesContext &ctx, const NodeProto &node, const char *x, const char *I,
                           const char *output_shape);

/**
 * Computes the output :cpp:class:`SymTensor` of a ``MaxRoiPool`` node and
 * stores it in ``ctx``.
 *
 * The output dtype matches the input feature-map dtype. The output shape
 * is ``(num_rois, C, pooled_shape[0], pooled_shape[1])`` where ``num_rois``
 * is taken from dim 0 of ``rois`` (which may be symbolic); ``C`` is taken
 * from dim 1 of ``x``; and the spatial sizes come from the required
 * ``pooled_shape`` attribute.
 *
 * @param ctx   In/out context. Must already contain entries for ``x`` and
 *              ``rois``; on return it also contains an entry for
 *              ``node.output(0)``.
 * @param node  The ``MaxRoiPool`` ``NodeProto`` whose output should be
 *              described. ``node.op_type()`` must be ``"MaxRoiPool"`` and
 *              ``node`` must declare at least one output.
 * @param x     Name of the feature-map input value (rank 4) in ``ctx``.
 * @param rois  Name of the RoIs input value (rank 2 with last dim 5) in
 *              ``ctx``.
 *
 * @throws std::invalid_argument if ``node.op_type()`` is not ``"MaxRoiPool"``,
 *         if ``node`` has no output, if any input has the wrong rank, or
 *         if ``pooled_shape`` is missing, has the wrong size, or contains
 *         non-positive values.
 * @throws std::out_of_range     if any input name is not present in ``ctx``.
 */
void ComputeShapeMaxRoiPool(ShapesContext &ctx, const NodeProto &node, const char *x,
                            const char *rois);

} // namespace nn
} // namespace ONNX_LIGHT_NAMESPACE::onnx_shapes::shapes
