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

#pragma once

#include "onnx_core/runtime/kernels/kernel_context.h"
#include "onnx_core/runtime/memory/simple_tensor.h"
#include "onnx_core/runtime/runtime_context.h"

#include <optional>
#include <span>
#include <string>
#include <vector>

namespace ONNX_LIGHT_NAMESPACE::onnx_kernels {
// Re-exports the runtime types moved to ``onnx_core::runtime`` so
// kernel implementations below can keep referring to them
// unqualified, matching pre-move code.
using namespace ::onnx_light::core::runtime;

namespace kernel {
using ::onnx_light::core::runtime::DefaultOpset;
using ::onnx_light::core::runtime::KernelBase;
using ::onnx_light::core::runtime::KernelContext;
using ::onnx_light::core::runtime::OpsetId;

// ---------------------------------------------------------------------------
// Reference implementations of the ``tensor`` backend test kernels.
//
// Each kernel is exposed as a small class whose constructor takes a
// :ref:`KernelContext` (carrying the opset against which the kernel must
// behave) and whose ``operator()`` performs the computation.
//
// Two flavors of ``operator()`` are provided:
//
//   * The returning overload (``Tensor operator()(...) const``) allocates a
//     fresh ``Tensor`` whose data buffer is owned by the returned value.
//   * The in-place overload (``void operator()(..., Tensor &output) const``)
//     writes results into a caller-supplied output tensor whose buffer has
//     already been allocated. The caller is responsible for setting
//     ``output.data_type``, ``output.shape`` and sizing ``output.data`` to
//     match the operator's expected output; the kernel validates these
//     attributes and throws ``std::invalid_argument`` on mismatch.
//
// Each kernel class also exposes a ``static constexpr bool CanRunInPlace()``
// query indicating whether the output tensor's data buffer may alias one of
// the input tensors' buffers. ``Concat`` returns ``false`` because its
// output is strictly larger (along ``axis``) than any single input and
// therefore cannot share storage with one.
// ---------------------------------------------------------------------------

/// Concatenates a list of tensors along ``axis`` (since opset 13). All input
/// tensors must share the same data type and the same shape except along the
/// concatenation axis. ``axis`` may be negative, in which case it counts from
/// the back of the input rank.
class Concat : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Concat";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensors &inputs, int64_t axis, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensors &inputs, int64_t axis, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Performs element-wise type conversion of an input tensor ``x`` to the
/// data type specified by ``to`` (a ``DataType`` value,
/// mirroring the ``Cast`` operator's required ``to`` attribute). The output
/// shape always matches the input shape.
///
/// The reference implementation supports the numeric element types in the
/// backend test library — ``FLOAT``, ``DOUBLE``, ``INT32``, ``INT64``,
/// ``INT8``, ``UINT8``, ``INT16``, ``UINT16`` and ``BOOL`` — as well as
/// ``STRING`` in either direction (numeric ↔ STRING uses the canonical
/// decimal representation). Other dtypes will cause the kernel to throw
/// ``std::invalid_argument``: this is sufficient for the backend test
/// cases registered today and keeps the implementation small.
/// Out-of-range floating-point values when casting to an integer dtype
/// follow C++ ``static_cast`` semantics, which matches the behaviour
/// exercised by the upstream ``test_cast_FLOAT_to_*`` node tests for the
/// supported conversions.
class Cast : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Cast";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, int32_t to, RuntimeContext *rt = nullptr) const;
  Tensor operator()(const Tensor &x, int32_t to, bool saturate, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, int32_t to, Tensor &output) const;
  void operator()(const Tensor &x, int32_t to, bool saturate, Tensor &output) const;

  /// Output element type may differ from the input element type, so storage
  /// can not be shared in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``BitCast`` operator (since
/// opset 26). Reinterprets the bit pattern of the input tensor as the
/// data type ``to`` without value conversion. ``to`` must be a non-string
/// type with the same element bit-width as ``x.data_type``; otherwise the
/// kernel throws ``std::invalid_argument``. Implementations treat the
/// underlying bytes as little endian, which matches the host ABIs
/// targeted by the backend test library.
class BitCast : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:BitCast";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, int32_t to, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, int32_t to, Tensor &output) const;

  /// In-place execution is permitted only when ``to == x.data_type``; the
  /// kernel itself does not enforce aliasing constraints so this flag is
  /// conservatively ``false``.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``AffineGrid`` operator (since
/// opset 20 in the ``ai.onnx`` domain). Generates a flow field of
/// sampling coordinates by applying a batch of affine matrices ``theta``
/// to a regular grid of size ``size``.
///
/// Inputs:
///   * ``theta``: FLOAT tensor of shape ``(N, 2, 3)`` for 2D or
///     ``(N, 3, 4)`` for 3D.
///   * ``size``: INT64 1-D tensor of length 4 (``(N, C, H, W)``) for 2D
///     or 5 (``(N, C, D, H, W)``) for 3D. Only the spatial dimensions
///     ``(H, W)`` or ``(D, H, W)`` are used; ``N`` is taken from
///     ``theta`` (and must match ``size[0]``) and ``C`` is ignored.
///
/// Attribute ``align_corners`` (int, default 0): when 1, the normalised
/// coordinates ``-1`` and ``+1`` refer to the centres of the corner
/// pixels; when 0 they refer to the outer edges (the convention matching
/// ``torch.nn.functional.affine_grid``).
///
/// Output shape: ``(N, H, W, 2)`` for 2D or ``(N, D, H, W, 3)`` for 3D.
/// The element type follows the ``theta`` input (FLOAT in this
/// implementation).
class AffineGrid : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:AffineGrid";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``AffineGrid`` operator.
  struct Attributes {
    int64_t align_corners = 0;
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &theta, const Tensor &size, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &theta, const Tensor &size, const Attributes &attrs, Tensor &output,
                  RawBufferAllocator *allocator = nullptr) const;

  /// Output shape and element layout differ from both inputs, so the
  /// output cannot share storage with any input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``GridSample`` operator (since
/// opset 16 in the ``ai.onnx`` domain; extended to N-D in opset 20).
/// Performs sampling of an input tensor ``X`` at the positions given by
/// the flow field ``grid``.
///
/// Inputs:
///   * ``X``: tensor of shape ``(N, C, D1, D2, ..., Dr)`` with ``r`` spatial
///     dimensions (``r >= 1``).
///   * ``grid``: floating-point tensor of shape
///     ``(N, D1_out, D2_out, ..., Dr_out, r)`` carrying normalised
///     sampling coordinates.
///
/// Attributes:
///   * ``mode`` (string, default ``"linear"`` / ``"bilinear"``):
///     interpolation mode, one of ``"linear"``/``"bilinear"``,
///     ``"nearest"`` or ``"cubic"``/``"bicubic"``.
///   * ``padding_mode`` (string, default ``"zeros"``): one of ``"zeros"``,
///     ``"border"`` or ``"reflection"``.
///   * ``align_corners`` (int, default 0): when 1, the normalised
///     coordinates ``-1`` and ``+1`` refer to the centres of the corner
///     pixels; when 0 they refer to the outer edges.
///
/// Output shape: ``(N, C, D1_out, D2_out, ..., Dr_out)``. The element
/// type follows ``X``; this implementation supports the FLOAT/DOUBLE
/// element types of ``X`` and ``grid``.
class GridSample : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:GridSample";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``GridSample`` operator.
  struct Attributes {
    std::string mode = "linear";
    std::string padding_mode = "zeros";
    int64_t align_corners = 0;
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &X, const Tensor &grid, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &X, const Tensor &grid, const Attributes &attrs, Tensor &output,
                  RawBufferAllocator *allocator = nullptr) const;

  /// Output shape differs from both inputs, so the output cannot share
  /// storage with any input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Performs element-wise type conversion of an input tensor ``x`` to the
/// data type carried by ``target_type`` (a second tensor whose values are
/// ignored). This mirrors the ONNX ``CastLike`` operator (since opset 15
/// in the ai.onnx domain), which is equivalent to ``Cast`` with
/// ``to = target_type.data_type``.
///
/// The reference implementation forwards to :ref:`kernel::Cast` and so
/// supports the same element-type matrix.
class CastLike : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:CastLike";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, const Tensor &target_type, RuntimeContext *rt = nullptr) const;
  Tensor operator()(const Tensor &x, const Tensor &target_type, bool saturate,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &target_type, Tensor &output) const;
  void operator()(const Tensor &x, const Tensor &target_type, bool saturate, Tensor &output) const;

  /// Output element type may differ from the input element type, so storage
  /// can not be shared in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Broadcasts the ``input`` tensor to the shape given by the 1-D INT64
/// ``shape`` tensor, following the ONNX numpy-style broadcasting rules
/// (ONNX ``Expand`` operator, since opset 8 in the ``ai.onnx`` domain).
///
/// The output shape is computed as ``broadcast(input.shape, shape_values)``.
/// A dimension in ``input`` of size 1 is expanded (repeated) to match the
/// corresponding target dimension; a dimension equal to the target is
/// left unchanged. The output dtype always matches the input dtype.
///
/// The reference implementation supports all whole-byte element types
/// supported by :cpp:func:`ElementSize`. String and sub-byte dtypes
/// (INT4/UINT4/INT2/UINT2) are not supported and will cause the kernel
/// to throw ``std::invalid_argument``.
class Expand : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Expand";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &input, const Tensor &shape, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, const Tensor &shape, Tensor &output) const;

  /// The output may be larger than either input, so storage cannot be
  /// shared in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reshapes ``data`` to the target shape described by the 1-D INT64 ``shape``
/// tensor (ONNX ``Reshape`` operator, since opset 5; with ``allowzero`` input
/// semantics unchanged in newer opsets).
///
/// Output dtype always matches ``data``. The output shape follows ONNX rules:
/// positive values are copied, one ``-1`` is inferred from element count, and
/// ``0`` copies the corresponding input dim unless ``allowzero`` is set.
class Reshape : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Reshape";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &shape, int64_t allowzero = 0,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &shape, int64_t allowzero, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Slices ``data`` according to ONNX ``Slice`` semantics (since opset 10+):
/// ``starts`` and ``ends`` are required; ``axes`` and ``steps`` are optional.
///
/// Supports positive and negative steps, negative indices, omitted axes/steps,
/// and clamping behavior aligned with ONNX shape-inference rules.
class Slice : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Slice";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &starts, const Tensor &ends,
                    const Tensor *axes = nullptr, const Tensor *steps = nullptr,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &starts, const Tensor &ends, const Tensor *axes,
                  const Tensor *steps, Tensor &output,
                  RawBufferAllocator *allocator = nullptr) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Pads the input ``data`` tensor along selected axes (ONNX ``Pad`` operator,
/// since opset 18 with the ``axes`` input; earlier semantics covered by
/// passing ``axes == nullptr``).
///
/// ``pads`` is a 1-D INT64 tensor of length ``2 * num_axes`` formatted as
/// ``[x1_begin, x2_begin, ..., x1_end, x2_end, ...]``. When ``axes`` is
/// ``nullptr`` it defaults to ``[0, 1, ..., rank-1]``. ``constant_value`` is
/// an optional scalar (one-element tensor) used only when ``mode`` is
/// ``"constant"`` (its dtype must match ``data``); when ``nullptr`` and the
/// mode is ``"constant"``, the kernel uses a zero-initialized value of the
/// input dtype. Supported modes mirror ``numpy.pad``: ``"constant"`` (default),
/// ``"reflect"``, ``"edge"`` and ``"wrap"`` (introduced in opset 19).
///
/// All padding entries must be non-negative (negative pads, i.e. cropping,
/// are accepted by the schema but the reference implementation rejects them
/// to keep the kernel focused on the cases covered by the bundled tests).
/// String and sub-byte dtypes are not supported and cause the kernel to
/// throw ``std::invalid_argument``.
class Pad : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Pad";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &pads, const Tensor *constant_value = nullptr,
                    const Tensor *axes = nullptr, const std::string &mode = "constant",
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &pads, const Tensor *constant_value,
                  const Tensor *axes, const std::string &mode, Tensor &output) const;

  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Constructs a tensor by tiling the ``input`` tensor a number of times along
/// each axis given by the 1-D INT64 ``repeats`` tensor (ONNX ``Tile``
/// operator, since opset 6 in the ``ai.onnx`` domain).
///
/// ``repeats`` must have the same length as ``input``'s rank, and every entry
/// must be non-negative. The output has the same rank and dtype as ``input``;
/// its dimension ``i`` is ``input.shape[i] * repeats[i]``.
///
/// The reference implementation supports all whole-byte element types
/// supported by :cpp:func:`ElementSize`. String and sub-byte dtypes
/// (INT4/UINT4/INT2/UINT2) are not supported and will cause the kernel to
/// throw ``std::invalid_argument``.
class Tile : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Tile";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &input, const Tensor &repeats, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, const Tensor &repeats, Tensor &output) const;

  /// The output is generally larger than the input, so storage cannot be
  /// shared in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Splits ``input`` along ``axis`` into a list of tensors (ONNX ``Split``
/// operator). When ``split`` is empty the input dimension on ``axis`` is
/// divided into ``num_outputs`` chunks of equal size (the last chunk being
/// smaller when the dimension is not evenly divisible). When ``split`` is
/// non-empty its entries give the size of each output along ``axis`` and
/// must sum to ``input.shape[axis]``. ``axis`` may be negative, in which
/// case it counts from the back of ``input``'s rank.
///
/// The reference implementation supports all whole-byte element types
/// supported by :cpp:func:`ElementSize`. STRING and sub-byte element types
/// are not supported and will cause the kernel to throw
/// ``std::invalid_argument``.
class Split : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Split";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  /// Computes the split. Exactly one of ``split`` (non-empty) or
  /// ``num_outputs`` (> 0) must be provided.
  ///
  /// This overload accepts a zero-copy span view of the split sizes (preferred
  /// when the sizes come directly from a tensor buffer).
  Tensors operator()(const Tensor &input, int64_t axis, std::span<const int64_t> split,
                     int64_t num_outputs, RuntimeContext *rt = nullptr) const;

  /// Convenience overload for brace-enclosed initializer lists such as
  /// ``{2, 4}``; copies the values into a temporary vector and forwards to
  /// the span overload.  Callers with an existing ``std::vector<int64_t>``
  /// can pass it directly via implicit conversion to ``std::span``.
  Tensors operator()(const Tensor &input, int64_t axis, std::initializer_list<int64_t> split,
                     int64_t num_outputs, RuntimeContext *rt = nullptr) const {
    std::vector<int64_t> tmp(split);
    return (*this)(input, axis, std::span<const int64_t>{tmp}, num_outputs, rt);
  }

  /// Output buffers are strict subsets of the input and have a different
  /// shape, so storage can not generally be shared.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Removes dimensions of size 1 from ``data`` according to ``axes`` (ONNX
/// ``Squeeze`` operator). When ``axes`` is empty, all dimensions with size 1
/// are removed.
class Squeeze : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Squeeze";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const onnx_kernels::Shape &axes,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const onnx_kernels::Shape &axes, Tensor &output) const;

  /// Rank may change after squeezing.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Inserts dimensions of size 1 into ``data`` at positions given by ``axes``
/// (ONNX ``Unsqueeze`` operator).
class Unsqueeze : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Unsqueeze";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const onnx_kernels::Shape &axes,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const onnx_kernels::Shape &axes, Tensor &output) const;

  /// Rank changes after unsqueezing.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Permutes the axes of the input tensor according to ``perm`` (ONNX
/// ``Transpose`` operator). When ``perm`` is empty, the axis order is
/// reversed.
///
/// The reference implementation supports whole-byte tensor element types
/// supported by :cpp:func:`ElementSize`.
class Transpose : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Transpose";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const onnx_kernels::Shape &perm,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const onnx_kernels::Shape &perm, Tensor &output) const;

  /// Output shape differs from input shape in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Returns the indices of the elements that are non-zero of the input tensor
/// ``X`` in row-major order, as a 2-D ``INT64`` tensor of shape ``(rank, nnz)``
/// (ONNX ``NonZero`` operator, since opset 9 in the ``ai.onnx`` domain). For
/// scalar input the output shape is ``(0, nnz)`` (mirroring the upstream
/// specification, which differs from NumPy).
///
/// The reference implementation supports the numeric and ``BOOL`` element
/// types in the backend test library — ``FLOAT``, ``DOUBLE``, ``INT8``,
/// ``UINT8``, ``INT16``, ``UINT16``, ``INT32``, ``INT64``, ``UINT32``,
/// ``UINT64`` and ``BOOL``. Other dtypes will cause the kernel to throw
/// ``std::invalid_argument``.
class NonZero : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:NonZero";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;

  /// The output has a different dtype (INT64) and a different shape
  /// from the input, so storage cannot be shared.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``OneHot`` operator (since opset 9 in
/// the ``ai.onnx`` domain; opset 11 added support for negative indices).
///
/// Inputs:
///   * ``indices``: rank ``>= 1`` numeric tensor of indices. Negative entries
///     are accepted from opset 11 onwards (negative ``i`` selects depth
///     position ``i + depth``); entries outside ``[-depth, depth-1]`` (or
///     ``[0, depth)`` for opset 9) leave the corresponding row filled with
///     ``off_value``.
///   * ``depth``: numeric scalar (or rank-1 tensor with a single element)
///     giving the size of the new one-hot dimension.
///   * ``values``: rank-1 tensor of length 2, formatted as
///     ``[off_value, on_value]``.
///
/// Attribute ``axis`` (int, default ``-1``): position at which the new
/// dimension is inserted in the output. Accepted range is
/// ``[-rank(indices)-1, rank(indices)]``.
///
/// The output has rank ``rank(indices) + 1``; its element type matches
/// ``values``. The reference implementation supports all numeric element
/// types for ``indices`` and ``depth`` (cast to ``int64`` before use) and
/// all numeric and ``BOOL`` element types for ``values`` / the output.
class OneHot : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:OneHot";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``OneHot`` operator.
  struct Attributes {
    int64_t axis = -1;
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &indices, const Tensor &depth, const Tensor &values,
                    const Attributes &attrs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &indices, const Tensor &depth, const Tensor &values,
                  const Attributes &attrs, Tensor &output) const;

  /// The output has a different shape (one extra dimension) than any of the
  /// inputs, so storage cannot be shared.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``Shape`` operator (since opset 1 in
/// the ``ai.onnx`` domain; extended with ``start``/``end`` attributes in
/// opset 15). Returns the shape of the input tensor as an ``INT64`` 1-D
/// tensor.
///
/// Attributes ``start`` (int, default 0) and ``end`` (int, optional) bound
/// the slice ``input.shape[start:end]`` (using numpy-style indexing).
/// Negative values count from the back; out-of-range values are clamped to
/// ``[0, r]`` where ``r`` is the rank of the input. When ``start > end``
/// (after normalisation) the output is empty.
///
/// The kernel reads only the input shape, never its data buffer, so it
/// accepts an input of any element type.
class Shape : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Shape";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``Shape`` operator.
  struct Attributes {
    int64_t start = 0;
    /// ``std::nullopt`` means "use the input rank" (no slicing on the right).
    std::optional<int64_t> end = std::nullopt;
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &data, RuntimeContext *rt = nullptr) const;
  Tensor operator()(const Tensor &data, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Attributes &attrs, Tensor &output) const;

  /// Output has a different dtype (INT64) and shape from the input, so
  /// storage cannot be shared.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``Size`` operator (since opset 1 in
/// the ``ai.onnx`` domain). Returns a 0-D (scalar) INT64 tensor whose single
/// value is the total number of elements of the input tensor, i.e. the
/// product of its dimensions.
///
/// The kernel reads only the input shape, never its data buffer, so it
/// accepts an input of any element type.
class Size : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Size";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &data, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, Tensor &output) const;

  /// Output has a different dtype (INT64) and shape (scalar) from the input,
  /// so storage cannot be shared.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``Identity`` operator (since opset 1
/// in the ``ai.onnx`` domain). Copies the input tensor to the output
/// unchanged. The output dtype and shape always match the input.
///
/// The kernel is element-type agnostic: it simply copies the raw data buffer
/// (and forwards string elements when applicable), so it accepts any tensor
/// element type supported by ``Tensor``.
class Identity : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Identity";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &input, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, Tensor &output) const;

  /// Output is a bit-for-bit copy of the input; storage may safely be shared
  /// when the caller decides so.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};
/// Reference implementation of the ONNX ``Gather`` operator (since opset 1 in
/// the ``ai.onnx`` domain). Gathers entries of the ``axis`` dimension of
/// ``data`` indexed by ``indices``, producing an output tensor of rank
/// ``q + (r - 1)`` where ``r = rank(data)`` and ``q = rank(indices)``.
///
/// ``indices`` may be INT32 or INT64; negative values count from the back of
/// the gathered axis. The output dtype always matches ``data``.
class Gather : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Gather";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &indices, int64_t axis = 0,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &indices, int64_t axis, Tensor &output,
                  RuntimeContext *rt = nullptr) const;

  /// Output shape differs from input shape in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``GatherElements`` operator (since
/// opset 11 in the ``ai.onnx`` domain). ``data`` and ``indices`` must have the
/// same rank ``r`` and the output has the same shape as ``indices``.
///
/// In the 3-D case: ``out[i][j][k] = data[indices[i][j][k]][j][k]`` when
/// ``axis == 0`` (and analogously for other axes). ``indices`` may be INT32 or
/// INT64; negative values count from the back of the gathered axis.
class GatherElements : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:GatherElements";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &indices, int64_t axis = 0,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &indices, int64_t axis, Tensor &output,
                  RuntimeContext *rt = nullptr) const;

  /// Output shape matches ``indices`` and differs from ``data`` in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``GatherND`` operator (since opset 11
/// in the ``ai.onnx`` domain; ``batch_dims`` attribute added in opset 12).
/// Gathers slices from ``data`` at the index tuples encoded by ``indices`` and
/// produces an output of rank ``q + r - indices_shape[-1] - 1 - b`` where
/// ``q = rank(indices)``, ``r = rank(data)`` and ``b = batch_dims``.
///
/// ``indices`` must be INT64. Negative index values count from the back of the
/// corresponding ``data`` axis.
class GatherND : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:GatherND";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &indices, int64_t batch_dims = 0,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &indices, int64_t batch_dims,
                  Tensor &output) const;

  /// Output shape differs from both inputs in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``TensorScatter`` operator (since
/// opset 24 in the ``ai.onnx`` domain). Performs a functional update of a
/// KV-cache style tensor by writing slices of ``update`` into a copy of
/// ``past_cache`` along the ``axis`` (sequence) dimension, starting at the
/// per-batch positions given by ``write_indices``.
///
/// Inputs:
///   * ``past_cache``: tensor of rank ``r >= 2`` with shape
///     ``(batch_size, D1, ..., max_sequence_length, ..., Dn)``.
///   * ``update``: tensor of the same dtype and rank as ``past_cache``;
///     the only dimension that may differ is the ``axis`` (sequence)
///     dimension which holds ``sequence_length``.
///   * ``write_indices``: optional 1-D ``INT64`` tensor of length
///     ``batch_size``. When omitted (``nullptr``) it defaults to all zeros.
///
/// Attributes ``axis`` (int, default ``-2``) and ``mode`` (string,
/// default ``"linear"``): ``axis`` selects the sequence dimension and
/// cannot be the batch dimension (``0``). When ``mode`` is ``"linear"``
/// every ``write_indices[batch] + sequence_idx`` must remain in
/// ``[0, max_sequence_length)``; when ``mode`` is ``"circular"`` the
/// write index wraps around modulo ``max_sequence_length``.
///
/// Output dtype and shape always match ``past_cache``.
class TensorScatter : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:TensorScatter";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``TensorScatter`` operator.
  struct Attributes {
    int64_t axis = -2;
    std::string mode = "linear";
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &past_cache, const Tensor &update, const Tensor *write_indices,
                    const Attributes &attrs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &past_cache, const Tensor &update, const Tensor *write_indices,
                  const Attributes &attrs, Tensor &output,
                  RawBufferAllocator *allocator = nullptr) const;

  /// Output shape matches ``past_cache`` so the output buffer could share
  /// storage with the first input; the reference implementation always
  /// writes into a freshly allocated buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Reference implementation of the ONNX ``ScatterElements`` operator (since
/// opset 11 in the ``ai.onnx`` domain; ``reduction`` attribute added in opset
/// 16 (``add``/``mul``) and extended in opset 18 (``max``/``min``)). ``data``,
/// ``indices`` and ``updates`` must have the same rank ``r`` and the output
/// has the same shape and dtype as ``data``. ``indices`` may be INT32 or
/// INT64; negative values count from the back of the scattered axis.
///
/// When ``reduction`` is ``"none"``, repeated indices result in unspecified
/// behaviour (the order of iteration is not defined). When ``reduction`` is
/// ``"add"``, ``"mul"``, ``"max"`` or ``"min"``, the corresponding element of
/// ``data`` is updated with ``f(out_elem, updates_elem)``.
class ScatterElements : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ScatterElements";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``ScatterElements`` operator.
  struct Attributes {
    int64_t axis = 0;
    std::string reduction = "none";
  };

  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &indices, const Tensor &updates,
                    const Attributes &attrs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &indices, const Tensor &updates,
                  const Attributes &attrs, Tensor &output, RuntimeContext *rt = nullptr) const;

  /// Output shape matches ``data`` so the output buffer could share storage
  /// with the first input; the reference implementation always writes into a
  /// freshly allocated buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Reference implementation of the deprecated ONNX ``Scatter`` operator
/// (opset 9, deprecated since opset 11 in favour of ``ScatterElements``).
/// Semantically equivalent to ``ScatterElements`` with ``reduction="none"``:
/// ``data``, ``indices`` and ``updates`` must have the same rank ``r`` and
/// the output has the same shape and dtype as ``data``. ``indices`` may be
/// INT32 or INT64; negative values count from the back of the scattered
/// axis.
class Scatter : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Scatter";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``Scatter`` operator.
  struct Attributes {
    int64_t axis = 0;
  };

  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &indices, const Tensor &updates,
                    const Attributes &attrs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &indices, const Tensor &updates,
                  const Attributes &attrs, Tensor &output) const;

  /// Output shape matches ``data`` so the output buffer could share storage
  /// with the first input; the reference implementation always writes into a
  /// freshly allocated buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Reference implementation of the ONNX ``ScatterND`` operator (since opset
/// 11 in the ``ai.onnx`` domain; ``reduction`` attribute added in opset 16
/// (``add``/``mul``) and extended in opset 18 (``max``/``min``)). The output
/// has the same shape and dtype as ``data``. ``indices`` must be INT64;
/// negative values count from the back of the corresponding ``data`` axis.
///
/// When ``reduction`` is ``"none"``, repeated indices result in unspecified
/// behaviour. When ``reduction`` is ``"add"``, ``"mul"``, ``"max"`` or
/// ``"min"``, the corresponding slice of ``data`` is updated with
/// ``f(out_slice, updates_slice)``.
class ScatterND : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ScatterND";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``ScatterND`` operator.
  struct Attributes {
    std::string reduction = "none";
  };

  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &data, const Tensor &indices, const Tensor &updates,
                    const Attributes &attrs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &data, const Tensor &indices, const Tensor &updates,
                  const Attributes &attrs, Tensor &output) const;

  /// Output shape matches ``data`` so the output buffer could share storage
  /// with the first input; the reference implementation always writes into a
  /// freshly allocated buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};
/// Reference implementation of the ONNX ``Compress`` operator (since opset 9
/// in the ``ai.onnx`` domain; axis may be negative since opset 11). Selects
/// slices from ``input`` along a given ``axis`` where the corresponding entry
/// of ``condition`` (a rank-1 BOOL tensor) is ``true``. When ``axis`` is not
/// supplied the input is first flattened and individual elements are selected;
/// the output is then a 1-D tensor.
///
/// The ``condition`` length may be shorter than the input size along the axis
/// (or the flattened size when no axis is given); excess slices are discarded.
///
/// The output dtype always matches ``input``. The selected count is a runtime
/// value and is therefore unknown at shape-inference time.
class Compress : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Compress";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  /// ``axis`` is an ``std::optional<int64_t>``: pass ``std::nullopt`` to
  /// compress the flattened input, or the axis index to compress along.
  Tensor operator()(const Tensor &input, const Tensor &condition, std::optional<int64_t> axis,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, const Tensor &condition, std::optional<int64_t> axis,
                  Tensor &output) const;

  /// Output size is data-dependent and cannot be inferred without evaluating
  /// ``condition`` at runtime.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``Trilu`` operator (since opset 14
/// in the ``ai.onnx`` domain). Returns the upper (``upper == 1``, default) or
/// lower (``upper == 0``) triangular part of the input tensor, keeping the
/// elements on/above (resp. on/below) the ``k``-th diagonal and zeroing the
/// others.
///
/// Inputs:
///   * ``input``: tensor of rank ``>= 2``. The last two dimensions
///     ``(N, M)`` are interpreted as a (batch of) matrix; any leading
///     dimensions are treated as batch dimensions and processed
///     independently.
///   * ``k``: optional ``INT64`` scalar (0-D tensor) selecting the
///     diagonal to keep/exclude. Defaults to ``0`` when omitted (signalled
///     by passing ``nullptr``).
///
/// Attribute ``upper`` (int, default 1): when 1 the upper triangular part
/// is retained; when 0 the lower triangular part is retained.
///
/// Output shape and dtype always match ``input``. Elements outside the
/// selected triangular region are set to zero (for ``BOOL`` to ``false``;
/// for ``STRING`` to the empty string).
class Trilu : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Trilu";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``Trilu`` operator.
  struct Attributes {
    int64_t upper = 1;
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &input, const Tensor *k, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, const Tensor *k, const Attributes &attrs,
                  Tensor &output) const;

  /// Output shape matches input shape so the output buffer could in theory
  /// share storage with the input; the reference implementation does not
  /// rely on that and always writes into a freshly allocated buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Reference implementation of the ONNX ``ReverseSequence`` operator (since
/// opset 10 in the ``ai.onnx`` domain). Reverses, for each batch slice ``i``,
/// the first ``sequence_lens[i]`` elements along the time axis. Elements past
/// that prefix are copied unchanged.
///
/// Inputs:
///   * ``input``: tensor of rank ``>= 2``. The ``time_axis`` and
///     ``batch_axis`` attributes (each one of ``0`` or ``1``) select which
///     of the first two dimensions plays the time / batch role. Any
///     remaining dimensions are treated as inner (feature) dimensions and
///     copied unchanged.
///   * ``sequence_lens``: rank-1 ``INT64`` tensor of length
///     ``input.shape[batch_axis]``. Each entry must satisfy
///     ``0 <= sequence_lens[i] <= input.shape[time_axis]``.
///
/// Attributes ``time_axis`` (int, default ``0``) and ``batch_axis``
/// (int, default ``1``): must be ``0`` or ``1`` and must differ.
///
/// Output shape and dtype always match ``input``.
class ReverseSequence : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ReverseSequence";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``ReverseSequence`` operator.
  struct Attributes {
    int64_t time_axis = 0;
    int64_t batch_axis = 1;
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &input, const Tensor &sequence_lens, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, const Tensor &sequence_lens, const Attributes &attrs,
                  Tensor &output) const;

  /// Output shape matches input shape so the kernel can in theory share
  /// storage with the input; the reference implementation always writes into
  /// a freshly allocated buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Reference implementation of the ONNX ``DepthToSpace`` operator (since opset
/// 1; ``mode`` attribute added in opset 11; type set extended in opset 13).
/// Rearranges (permutes) data from depth into blocks of spatial data — the
/// inverse of ``SpaceToDepth``. The input must be a 4-D tensor of shape
/// ``(N, C, H, W)`` with ``C`` divisible by ``blocksize * blocksize``. The
/// output has shape ``(N, C/(blocksize*blocksize), H*blocksize, W*blocksize)``.
///
/// ``mode`` is either ``"DCR"`` (default; depth-column-row order) or ``"CRD"``
/// (column-row-depth order); the reference implementation matches the
/// upstream NumPy equivalents in the operator spec.
///
/// The reference implementation supports whole-byte tensor element types
/// supported by :cpp:func:`ElementSize`.
class DepthToSpace : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:DepthToSpace";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``DepthToSpace`` operator.
  struct Attributes {
    int64_t blocksize = 0;
    std::string mode = "DCR";
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &input, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, const Attributes &attrs, Tensor &output) const;

  /// Output shape differs from input shape in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``SpaceToDepth`` operator (since opset
/// 1; type set extended in opset 13). Rearranges blocks of spatial data into
/// depth — the inverse of ``DepthToSpace``. The input must be a 4-D tensor of
/// shape ``(N, C, H, W)`` with ``H`` and ``W`` both divisible by
/// ``blocksize``. The output has shape
/// ``(N, C*blocksize*blocksize, H/blocksize, W/blocksize)``.
///
/// The reference implementation supports whole-byte tensor element types
/// supported by :cpp:func:`ElementSize`.
class SpaceToDepth : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:SpaceToDepth";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``SpaceToDepth`` operator.
  struct Attributes {
    int64_t blocksize = 0;
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &input, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, const Attributes &attrs, Tensor &output) const;

  /// Output shape differs from input shape in general.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``CenterCropPad`` operator (since
/// opset 18 in the ``ai.onnx`` domain). It centrally crops and/or pads the
/// input tensor to the dimensions given by the ``shape`` input. The
/// ``shape`` input is a 1-D tensor of length ``rank(input_data)`` when
/// ``axes`` is unset, or ``len(axes)`` otherwise. The ``axes`` attribute,
/// when set, specifies the subset of input axes ``shape`` refers to;
/// unspecified axes are left unchanged. Negative axes are supported.
///
/// On each selected axis the kernel matches the upstream reference: when
/// ``shape[i] < input_dim`` a centered cropping window of length
/// ``shape[i]`` is extracted (the start index is rounded down on odd
/// differences); when ``shape[i] > input_dim`` the input is zero-padded
/// equally on both sides (an extra pixel is added on the right when the
/// total padding is odd).
class CenterCropPad : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:CenterCropPad";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``CenterCropPad`` operator. ``axes`` is
  /// optional; an empty vector here means "all axes" (i.e. ``shape`` has
  /// length equal to ``rank(input_data)``).
  struct Attributes {
    onnx_kernels::Shape axes;
    bool axes_present = false;
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &input_data, const Tensor &shape, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input_data, const Tensor &shape, const Attributes &attrs,
                  Tensor &output) const;

  /// Output shape generally differs from input shape, so storage cannot be
  /// shared with the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the (deprecated) ONNX ``Upsample`` operator
/// (since opset 1 in the ``ai.onnx`` domain, last refreshed at opset 10 and
/// replaced by ``Resize`` from opset 10 onwards).
///
/// The kernel takes the input tensor ``X`` and a 1-D ``FLOAT`` ``scales``
/// tensor of length ``rank(X)`` and produces a tensor whose dim ``i`` is
/// ``floor(X.shape[i] * scales[i])``. The ``mode`` attribute selects the
/// interpolation rule:
///
///   * ``"nearest"`` (default) — every output element is copied from the
///     nearest input element using ``floor(out / scale)`` to map output
///     coordinates back to input coordinates. Supports any input rank.
///   * ``"linear"`` / ``"bilinear"`` — supported only for 4-D NCHW input,
///     with scales equal to ``1`` on the ``N`` and ``C`` axes. The two
///     spatial axes are upsampled with the "asymmetric" coordinate
///     transformation used by Upsample v7/9/10 in upstream ONNX.
///
/// The reference implementation supports the same whole-byte element
/// types as :cpp:func:`ElementSize` for ``"nearest"`` mode; ``"linear"``
/// mode requires a floating-point input (``FLOAT16``/``FLOAT``/``DOUBLE``).
class Upsample : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Upsample";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``Upsample`` operator.
  struct Attributes {
    std::string mode = "nearest";
  };

  using KernelBase::KernelBase;

  Tensor operator()(const Tensor &X, const Tensor &scales, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &X, const Tensor &scales, const Attributes &attrs,
                  Tensor &output) const;

  /// Output shape generally differs from input shape, so storage cannot be
  /// shared with the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``Resize`` operator (since opset 10
/// in the ``ai.onnx`` domain, replacing the deprecated ``Upsample``).
///
/// The kernel computes ``Y`` from input ``X`` and either ``scales`` (a 1-D
/// FLOAT tensor with one entry per input axis) or ``sizes`` (a 1-D INT64
/// tensor giving the target shape). Exactly one of ``scales`` and ``sizes``
/// must be provided. The output dim ``i`` is ``floor(X.shape[i] *
/// scales[i])`` when ``scales`` is used and ``sizes[i]`` when ``sizes`` is
/// used.
///
/// The ``"nearest"``, ``"linear"`` and ``"cubic"`` modes of the ``mode``
/// attribute are supported. The input coordinate of an output position is
/// computed according to ``coordinate_transformation_mode`` (``half_pixel``
/// -- the default -- ``half_pixel_symmetric``, ``pytorch_half_pixel``,
/// ``align_corners`` and ``asymmetric``). In ``"nearest"`` mode it is
/// rounded according to ``nearest_mode`` and clamped to ``[0, in_dim - 1]``;
/// all four ONNX ``nearest_mode`` values (``round_prefer_floor`` -- the
/// default -- ``round_prefer_ceil``, ``floor``, ``ceil``) are supported.
/// In ``"linear"`` and ``"cubic"`` mode, the output is computed as a
/// separable per-axis weighted sum of neighbouring input samples (2 taps
/// for ``"linear"``, 4 taps for ``"cubic"``); the ``cubic_coeff_a``
/// attribute (default ``-0.75``) tunes the cubic kernel and
/// ``exclude_outside`` (default ``0``) renormalises the coefficients when
/// the support window falls outside the input range. The optional ``axes``
/// attribute (opset 18) selects the subset of axes ``scales``/``sizes``
/// refer to, and ``keep_aspect_ratio_policy`` (``"stretch"`` -- the default
/// -- ``"not_larger"``, ``"not_smaller"``) is honoured when ``sizes`` is
/// used. The ``"tf_crop_and_resize"`` coordinate transformation is
/// supported (using the ``roi`` argument provided through
/// :cpp:member:`Attributes::roi` and the ``extrapolation_value`` attribute);
/// ``antialias`` is supported for ``"linear"`` and ``"cubic"`` modes (when
/// non-zero, the interpolation kernel is widened by the resize scale while
/// downsampling to act as an anti-aliasing filter). The supported element
/// types are the same whole-byte types as :cpp:func:`ElementSize` for
/// ``"nearest"`` mode; ``"linear"`` and ``"cubic"`` modes require a
/// floating-point input (``FLOAT`` or ``DOUBLE``).
class Resize : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Resize";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``Resize`` operator.
  struct Attributes {
    std::string mode = "nearest";
    std::string coordinate_transformation_mode = "half_pixel";
    std::string nearest_mode = "round_prefer_floor";
    /// Subset of input axes that ``scales``/``sizes`` apply to. When empty,
    /// every axis is resized (matching the pre-opset-18 behaviour).
    onnx_kernels::Shape axes;
    /// How to interpret ``sizes``. ``"stretch"`` honours ``sizes`` exactly;
    /// ``"not_larger"`` / ``"not_smaller"`` rescale by a common factor so
    /// that the resulting output is not larger / smaller than ``sizes`` on
    /// any of the selected ``axes`` (other axes are left untouched).
    std::string keep_aspect_ratio_policy = "stretch";
    /// Coefficient ``A`` of the cubic interpolation kernel (used only when
    /// ``mode == "cubic"``). Defaults to ``-0.75`` to match the ONNX spec.
    float cubic_coeff_a = -0.75f;
    /// When non-zero (and ``mode`` is ``"linear"`` or ``"cubic"``),
    /// coefficients corresponding to neighbours that fall outside the
    /// input range are zeroed and the remaining coefficients are
    /// renormalised to sum to 1.
    int64_t exclude_outside = 0;
    /// Value used when ``coordinate_transformation_mode`` is
    /// ``"tf_crop_and_resize"`` and the transformed coordinate falls outside
    /// ``[0, input_dim - 1]``. Ignored for every other coordinate
    /// transformation mode.
    float extrapolation_value = 0.0f;
    /// Region of interest (used only by the ``"tf_crop_and_resize"``
    /// coordinate transformation). Stored as ``[start_axis_0, ...,
    /// start_axis_{N-1}, end_axis_0, ..., end_axis_{N-1}]``, where ``N`` is
    /// the number of resized axes (i.e. ``axes.size()`` when ``axes`` is
    /// non-empty, otherwise the input rank).
    std::vector<float> roi;
    /// When non-zero (and ``mode`` is ``"linear"`` or ``"cubic"``), the
    /// interpolation kernel is stretched by the resize scale when
    /// downsampling so that it acts as an anti-aliasing low-pass filter.
    /// Ignored for ``"nearest"`` mode and for upsampling (scale >= 1).
    int64_t antialias = 0;
  };

  using KernelBase::KernelBase;

  /// Resize using the ``scales`` input only. The ``sizes`` input is treated
  /// as absent (matching the ``Resize(X, scales)`` convenience form).
  Tensor operator()(const Tensor &X, const Tensor &scales, const Attributes &attrs,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &X, const Tensor &scales, const Attributes &attrs,
                  Tensor &output) const;

  /// Resize using the ``sizes`` input (1-D INT64 tensor giving the target
  /// shape). The ``scales`` input must be empty.
  Tensor ResizeSizes(const Tensor &X, const Tensor &sizes, const Attributes &attrs,
                     RuntimeContext *rt = nullptr) const;

  /// Output shape generally differs from input shape, so storage cannot be
  /// shared with the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``Unique`` operator (since opset 11
/// in the ``ai.onnx`` domain). Returns the unique values or subtensors of
/// the input tensor and three optional companion outputs:
///
///   * ``Y``               — unique values (1-D when ``axis`` is not provided)
///                           or unique subtensors sliced along ``axis``;
///   * ``indices``         — 1-D INT64 indices of the first occurrence of
///                           every ``Y`` element in ``X``;
///   * ``inverse_indices`` — 1-D INT64 indices that map every ``X`` element
///                           back to its position in ``Y``;
///   * ``counts``          — 1-D INT64 count of each ``Y`` element in ``X``.
///
/// When ``sorted`` is ``true`` (the default) outputs are ordered by ascending
/// value (or lexicographic order when ``axis`` is provided). When ``sorted``
/// is ``false`` the order of first occurrence in ``X`` is preserved.
///
/// The reference implementation supports the same whole-byte element types
/// as :cpp:func:`ElementSize` (FLOAT, DOUBLE, INT8/16/32/64, UINT8/16/32/64,
/// BOOL) and STRING.
class Unique : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Unique";
  void Run(RuntimeContext &rt) override;
  /// Attributes carried by the ONNX ``Unique`` operator.
  struct Attributes {
    /// Whether to sort unique elements in ascending order. Defaults to
    /// ``true`` (matching the schema default of 1).
    bool sorted = true;
    /// Axis along which to take unique subtensors. ``std::nullopt`` means
    /// "flatten the input" (matching the schema default behaviour).
    std::optional<int64_t> axis = std::nullopt;
  };

  /// Aggregated output of :cpp:class:`Unique`.
  struct Outputs {
    Tensor y;
    Tensor indices;
    Tensor inverse_indices;
    Tensor counts;
  };

  using KernelBase::KernelBase;

  Outputs operator()(const Tensor &x) const;
  Outputs operator()(const Tensor &x, const Attributes &attrs, RuntimeContext *rt = nullptr) const;

  /// The outputs have different dtypes (Y matches X; the others are INT64)
  /// and different shapes, so storage cannot be shared with the input.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

} // namespace kernel
} // namespace ONNX_LIGHT_NAMESPACE::onnx_kernels
