// 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 "onnx_extensions/kernels/tuning/portable_gemm_tuning.h"
#include "onnx_extensions/kernels/tuning/portable_parallel_tuning.h"

#include <memory>
#include <string>
#include <utility>
#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 ``math`` 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 kernel may be invoked with an output tensor
// whose data buffer aliases one of the input tensors' buffers (i.e. the
// caller is allowed to reuse an input buffer as the output buffer). The
// query returns the operator-level capability; concrete invocations must
// still verify that input and output shapes/types match. Operators whose
// output cannot share storage with any input (different element type, or
// larger output than every input — e.g. ``BlackmanWindow``, ``Concat``)
// return ``false``.
// ---------------------------------------------------------------------------

/// Element-wise absolute value.
class Abs : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Abs";
  explicit Abs(const KernelContext &ctx);
  /// Registers one portable tuning schema for every supported element type.
  static void RegisterTuningSchemas();
  KernelTuningKey TuningKey(int32_t element_type) const override;
  void Configure(const KernelTuningParameters &parameters) override;
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Returns the immutable configuration used by the execution path.
  const tuning::ParallelTuning &tuning() const noexcept { return tuning_; }

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }

private:
  tuning::ParallelTuning tuning_;
};

/// Element-wise arc cosine: y = acos(x), with x in [-1, 1] and y in [0, pi].
class Acos : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Acos";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise inverse hyperbolic cosine: y = acosh(x), with x >= 1.
class Acosh : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Acosh";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise arc sine: y = asin(x), with x in [-1, 1] and y in [-pi/2, pi/2].
class Asin : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Asin";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise inverse hyperbolic sine: y = asinh(x), defined for all real x.
class Asinh : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Asinh";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise arc tangent: y = atan(x), defined for all real x.
class Atan : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Atan";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise inverse hyperbolic tangent: y = atanh(x), with x in (-1, 1).
class Atanh : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Atanh";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise cosine: y = cos(x), defined for all real x with y in [-1, 1].
class Cos : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Cos";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise hyperbolic cosine: y = cosh(x), defined for all real x with y >= 1.
class Cosh : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Cosh";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Determinant of a square matrix or batches of square matrices.
///
/// The input must have shape ``[*, M, M]`` (rank >= 2) with the inner-most
/// two dimensions forming square matrices. The output has shape ``[*]``
/// (the leading batch dimensions); when the input is 2-D the output is a
/// scalar (rank-0). Both input and output are FLOAT tensors.
class Det : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Det";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output,
                  RawBufferAllocator *temporary_allocator = nullptr) const;

  /// The output buffer is smaller than the input (drops the trailing two
  /// dimensions); storage can never be shared with the input.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise ceiling: y = ceil(x), the smallest integer >= x.
class Ceil : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Ceil";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise floor: y = floor(x), the largest integer <= x.
class Floor : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Floor";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise clip: y = min(max(x, min_val), max_val) (since opset 6).
///
/// Both ``min`` and ``max`` are optional scalar tensors (since opset 11);
/// when omitted, the corresponding bound defaults to
/// ``std::numeric_limits<T>::lowest()`` and ``std::numeric_limits<T>::max()``.
/// When ``min > max`` the operator returns ``max`` for every element,
/// matching the ONNX specification (``Min(max, Max(input, min))``).
class Clip : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Clip";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  /// Computes ``y = clip(x, min, max)``. ``min``/``max`` may be ``nullptr``
  /// to use the dtype-specific default bound. When provided, each must be a
  /// 0-D (scalar) tensor whose dtype matches ``x``.
  Tensor operator()(const Tensor &x, const Tensor *min = nullptr, const Tensor *max = nullptr,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor *min, const Tensor *max, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise sign: y = sign(x), returning -1 for x<0, 0 for x==0, +1 for x>0.
class Sign : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Sign";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise round to nearest integer, ties to even (banker's rounding).
class Round : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Round";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise exponential: y = exp(x), defined for all real x.
class Exp : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Exp";
  explicit Exp(const KernelContext &ctx);
  /// Registers one portable tuning schema for every supported element type.
  static void RegisterTuningSchemas();
  KernelTuningKey TuningKey(int32_t element_type) const override;
  void Configure(const KernelTuningParameters &parameters) override;
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Returns the immutable configuration used by the execution path.
  const tuning::ParallelTuning &tuning() const noexcept { return tuning_; }

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }

private:
  tuning::ParallelTuning tuning_;
};

/// Element-wise error function: y = erf(x), defined for all real x.
class Erf : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Erf";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise natural logarithm: y = log(x), with x > 0.
class Log : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Log";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise reciprocal: y = 1 / x (Inf for x = 0).
class Reciprocal : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Reciprocal";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise square root: y = sqrt(x), with x >= 0 (NaN otherwise).
class Sqrt : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Sqrt";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise logistic sigmoid: y = 1 / (1 + exp(-x)).
class Sigmoid : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Sigmoid";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise HardSigmoid activation:
///     y = max(0, min(1, alpha * x + beta)).
class HardSigmoid : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:HardSigmoid";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, float alpha, float beta, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, float alpha, float beta, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise HardSwish activation:
///     y = x * max(0, min(1, x/6 + 0.5)) = x * HardSigmoid(x; alpha=1/6, beta=0.5).
class HardSwish : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:HardSwish";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Hardmax: emits 1 at the position of the first maximum value along ``axis``,
/// 0 elsewhere. The output has the same shape and dtype as the input.
class Hardmax : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Hardmax";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, int64_t axis, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, int64_t axis, Tensor &output) const;

  /// Hardmax must look at every value along the axis before it can write any
  /// output, so it cannot run in place safely.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise sine: y = sin(x), defined for all real x with y in [-1, 1].
class Sin : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Sin";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Softmax normalized exponential along a selected axis.
class Softmax : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Softmax";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, int64_t axis, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, int64_t axis, Tensor &output) const;

  /// Softmax needs the full input slice to compute each output value; aliasing
  /// input/output would overwrite values needed for later positions.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// LogSoftmax: log of the normalized exponential along a selected axis.
class LogSoftmax : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:LogSoftmax";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, int64_t axis, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, int64_t axis, Tensor &output) const;

  /// LogSoftmax needs the full input slice to compute each output value;
  /// aliasing input/output would overwrite values needed for later positions.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise softplus activation: y = ln(1 + exp(x)).
class Softplus : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Softplus";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise Mish activation: y = x * tanh(softplus(x)) =
/// x * tanh(ln(1 + exp(x))).
class Mish : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Mish";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise softsign activation: y = x / (1 + |x|).
class Softsign : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Softsign";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// SoftmaxCrossEntropyLoss computes the cross-entropy loss between the
/// (un-normalized) softmax distribution of ``scores`` and integer class
/// indices given by ``labels``. Optionally supports per-class weights and
/// an ``ignore_index`` value used to mask labels that should not contribute
/// to the loss. The kernel always returns the loss tensor and the
/// log-probability tensor (same shape and dtype as ``scores``); the caller
/// is free to ignore ``log_prob`` when the node does not request it.
class SoftmaxCrossEntropyLoss : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:SoftmaxCrossEntropyLoss";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  /// @param scores Input scores of shape ``(N, C)`` or ``(N, C, D1, ..., Dk)``.
  /// @param labels Integer class indices of shape ``(N)`` or ``(N, D1, ..., Dk)``.
  /// @param weights Optional rank-1 tensor of length ``C``; ``nullptr`` means unweighted.
  /// @param reduction One of ``"none"``, ``"sum"``, or ``"mean"``.
  /// @param has_ignore_index Whether ``ignore_index`` is set.
  /// @param ignore_index Class index to ignore (only used when ``has_ignore_index`` is true).
  /// @return Pair ``(loss, log_prob)`` with the loss tensor (shape per ``reduction``) and the
  ///         log-probability tensor (same shape and dtype as ``scores``).
  std::pair<Tensor, Tensor> operator()(const Tensor &scores, const Tensor &labels,
                                       const Tensor *weights, const std::string &reduction,
                                       bool has_ignore_index, int64_t ignore_index,
                                       RuntimeContext *rt = nullptr) const;

  /// The kernel allocates new outputs; it does not support input/output aliasing.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// NegativeLogLikelihoodLoss computes the (optionally weighted) negative
/// log-likelihood loss given an ``input`` tensor of log-probabilities of
/// shape ``(N, C)`` or ``(N, C, D1, ..., Dk)`` and integer class indices
/// ``target`` of shape ``(N)`` or ``(N, D1, ..., Dk)``. Unlike
/// ``SoftmaxCrossEntropyLoss``, the input is assumed to already contain
/// log-probabilities (no softmax is applied internally).
class NegativeLogLikelihoodLoss : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:NegativeLogLikelihoodLoss";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  /// @param input Log-probabilities of shape ``(N, C)`` or ``(N, C, D1, ..., Dk)``.
  /// @param target Integer class indices of shape ``(N)`` or ``(N, D1, ..., Dk)``.
  /// @param weight Optional rank-1 tensor of length ``C``; ``nullptr`` means unweighted.
  /// @param reduction One of ``"none"``, ``"sum"``, or ``"mean"``.
  /// @param has_ignore_index Whether ``ignore_index`` is set.
  /// @param ignore_index Class index to ignore (only used when ``has_ignore_index`` is true).
  /// @return Loss tensor whose shape depends on ``reduction``: matches the
  ///         shape of ``target`` for ``"none"``; scalar otherwise.
  Tensor operator()(const Tensor &input, const Tensor &target, const Tensor *weight,
                    const std::string &reduction, bool has_ignore_index, int64_t ignore_index,
                    RuntimeContext *rt = nullptr) const;

  /// The kernel allocates new outputs; it does not support input/output aliasing.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise hyperbolic sine: y = sinh(x), defined for all real x.
class Sinh : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Sinh";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise tangent: y = tan(x); undefined at x = (2k+1) * pi/2.
class Tan : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Tan";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise hyperbolic tangent: y = tanh(x), with y in (-1, 1).
class Tanh : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Tanh";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise rectified linear unit: ``y = max(0, x)``. Defined for both
/// floating-point and signed integer tensors (see schema since v14).
class Relu : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Relu";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise exponential linear unit:
/// ``y = x`` for ``x >= 0`` and ``y = alpha * (exp(x) - 1)`` otherwise.
/// ``alpha`` defaults to 1.0 to match the ONNX schema.
class Elu : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Elu";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, float alpha = 1.0f, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, float alpha, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise continuously differentiable exponential linear unit:
/// ``y = max(0, x) + min(0, alpha * (exp(x / alpha) - 1))``.
/// ``alpha`` defaults to 1.0 to match the ONNX schema. Only ``FLOAT`` inputs
/// are accepted (matching the v12 schema).
class Celu : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Celu";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, float alpha = 1.0f, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, float alpha, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise gaussian error linear unit. When ``approximate == "none"``
/// (the default) the exact formulation ``y = 0.5 * x * (1 + erf(x / sqrt(2)))``
/// is used; when ``approximate == "tanh"`` the tanh-based approximation
/// ``y = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`` is used.
class Gelu : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Gelu";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, const std::string &approximate = "none",
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const std::string &approximate, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise Swish activation: y = x * sigmoid(alpha * x). ``alpha``
/// defaults to 1.0 to match the ONNX schema (opset 24).
class Swish : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Swish";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, float alpha = 1.0f, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, float alpha, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Gated SwiGLU activation: ``Y = Swish_alpha(A) * B`` where
/// ``Swish_alpha(a) = a * sigmoid(alpha * a)``. Inputs ``A`` (gate) and ``B``
/// (value) must have identical shapes and dtypes: broadcasting is not applied.
/// ``alpha`` defaults to 1.0 to match the ONNX schema (opset 28).
class SwiGLU : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:SwiGLU";
  explicit SwiGLU(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &a, const Tensor &b, float alpha = 1.0f,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &a, const Tensor &b, float alpha, Tensor &output) const;
};

/// Element-wise thresholded rectified linear unit: y = x for x > alpha, y = 0
/// otherwise. ``alpha`` defaults to 1.0 to match the ONNX schema.
class ThresholdedRelu : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:ThresholdedRelu";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, float alpha = 1.0f, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, float alpha, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise leaky rectified linear unit:
/// ``y = x`` for ``x >= 0`` and ``y = alpha * x`` otherwise.
/// ``alpha`` defaults to 0.01 to match the ONNX schema.
class LeakyRelu : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:LeakyRelu";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, float alpha = 0.01f, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, float alpha, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise Shrink activation:
/// ``y = x + bias`` if ``x < -lambd``, ``y = x - bias`` if ``x > lambd``,
/// ``y = 0`` otherwise. ``lambd`` defaults to 0.5 and ``bias`` defaults to 0.
class Shrink : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Shrink";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, float bias = 0.0f, float lambd = 0.5f,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, float bias, float lambd, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise scaled exponential linear unit:
/// ``y = gamma * x`` for ``x > 0`` and
/// ``y = gamma * (alpha * exp(x) - alpha)`` for ``x <= 0``.
/// ``alpha`` and ``gamma`` default to the ONNX schema defaults
/// (~1.6732632 and ~1.0507009 respectively).
class Selu : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Selu";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, float alpha = 1.67326319217681884765625f,
                    float gamma = 1.05070102214813232421875f, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, float alpha, float gamma, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise addition with NumPy-style broadcasting.
class Add : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Add";
  explicit Add(const KernelContext &ctx);
  /// Registers one portable tuning schema for every supported element type.
  static void RegisterTuningSchemas();
  KernelTuningKey TuningKey(int32_t element_type) const override;
  void Configure(const KernelTuningParameters &parameters) override;
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  /// Returns the immutable configuration used by the execution path.
  const tuning::ParallelTuning &tuning() const noexcept { return tuning_; }

  /// Element-wise binary kernel: the output buffer may alias an input buffer
  /// when that input is not broadcast-expanded (i.e. its shape equals the
  /// output shape).
  static constexpr bool CanRunInPlace() noexcept { return true; }

private:
  tuning::ParallelTuning tuning_;
};

/// Element-wise subtraction with NumPy-style broadcasting.
class Sub : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Sub";
  explicit Sub(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  /// Element-wise binary kernel: the output buffer may alias an input buffer
  /// when that input is not broadcast-expanded (i.e. its shape equals the
  /// output shape).
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise multiplication with NumPy-style broadcasting.
class Mul : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Mul";
  explicit Mul(const KernelContext &ctx);
  /// Registers one portable tuning schema for every supported element type.
  static void RegisterTuningSchemas();
  KernelTuningKey TuningKey(int32_t element_type) const override;
  void Configure(const KernelTuningParameters &parameters) override;
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  /// Returns the immutable configuration used by the execution path.
  const tuning::ParallelTuning &tuning() const noexcept { return tuning_; }

  /// Element-wise binary kernel: the output buffer may alias an input buffer
  /// when that input is not broadcast-expanded (i.e. its shape equals the
  /// output shape).
  static constexpr bool CanRunInPlace() noexcept { return true; }

private:
  tuning::ParallelTuning tuning_;
};

/// Element-wise parametric ReLU: ``y = x`` when ``x >= 0`` and
/// ``y = slope * x`` otherwise. The ``slope`` tensor is unidirectionally
/// broadcastable to ``x`` (the standard NumPy-style multidirectional
/// broadcasting rules apply here as well — the unidirectional case is a
/// strict subset). Float-input semantics preserve ``+inf``/``-inf`` because
/// the kernel branches on the sign of ``x`` rather than evaluating
/// ``max(0, x) + slope * min(0, x)`` (which would yield ``NaN`` for
/// infinite inputs; see microsoft/onnxruntime#28732).
class PRelu : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:PRelu";
  explicit PRelu(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &slope, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &slope, Tensor &output) const;

  /// Element-wise binary kernel: the output buffer may alias an input buffer
  /// when that input is not broadcast-expanded (i.e. its shape equals the
  /// output shape).
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise division with NumPy-style broadcasting.
class Div : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Div";
  explicit Div(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  /// Element-wise binary kernel: the output buffer may alias an input buffer
  /// when that input is not broadcast-expanded (i.e. its shape equals the
  /// output shape).
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise modulo with NumPy-style broadcasting. The ``fmod`` flag
/// controls the semantics:
///   * ``fmod == 0`` (default): integer modulo whose sign follows the divisor
///     (Python ``%`` / ``numpy.mod``). Only valid for integer dtypes.
///   * ``fmod == 1``: C ``fmod`` semantics whose sign follows the dividend.
///     Required when either input is floating point and also accepted for
///     integer inputs (where it coincides with C ``%`` truncated modulo).
class Mod : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Mod";
  explicit Mod(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, int64_t fmod = 0,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, int64_t fmod, Tensor &output) const;

  /// Element-wise binary kernel: the output buffer may alias an input buffer
  /// when that input is not broadcast-expanded (i.e. its shape equals the
  /// output shape).
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise negation: y = -x. Mirrors the upstream ONNX ``Neg``
/// reference implementation (``np.negative``). Only the FLOAT dtype is
/// supported.
class Neg : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Neg";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, Tensor &output) const;

  /// Element-wise unary kernel: the output buffer may alias the input buffer.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise exponentiation ``z = x ^ y`` with NumPy-style multidirectional
/// broadcasting (since opset 7). Unlike most binary element-wise kernels,
/// ``Pow`` allows ``x`` and ``y`` to have different dtypes: the output dtype
/// always matches the dtype of the base ``x`` (per the ONNX schema, ``Z`` is
/// constrained to ``T`` while ``Y`` is constrained to the broader ``T1``).
///
/// Supported base dtypes (``T``): FLOAT, INT32, INT64.
/// Supported exponent dtypes (``T1``): FLOAT, INT32, INT64, UINT32, UINT64.
/// Integer base / integer exponent pairs evaluate the power in ``double``
/// precision and cast the result back to the base dtype, matching the
/// reference behaviour of NumPy's ``numpy.power`` and the upstream ONNX
/// backend test cases.
class Pow : public tuning::ParallelTunableKernel {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Pow";
  explicit Pow(const KernelContext &ctx);
  static void RegisterTuningSchemas();
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &x, const Tensor &y, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &y, Tensor &output) const;

  /// Element-wise binary kernel: the output buffer may alias the base input
  /// buffer when that input is not broadcast-expanded (i.e. its shape equals
  /// the output shape). Aliasing the exponent is generally not safe because
  /// it may have a different dtype than the output.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// BlackmanWindow function evaluated at ``size`` integer samples. When
/// ``periodic`` is true the window is computed as if of length ``size+1`` and
/// the last sample is discarded (matches NumPy/ONNX conventions).
class BlackmanWindow : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:BlackmanWindow";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &size, bool periodic = true, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &size, bool periodic, Tensor &output) const;

  /// Output is a float vector while the input is an int scalar: storage
  /// can never be shared with an input.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// HannWindow function evaluated at ``size`` integer samples. When
/// ``periodic`` is true the window is computed as if of length ``size+1`` and
/// the last sample is discarded (matches NumPy/ONNX conventions).
class HannWindow : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:HannWindow";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &size, bool periodic = true, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &size, bool periodic, Tensor &output) const;

  /// Output is a float vector while the input is an int scalar: storage
  /// can never be shared with an input.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// HammingWindow function evaluated at ``size`` integer samples. When
/// ``periodic`` is true the window is computed as if of length ``size+1`` and
/// the last sample is discarded (matches NumPy/ONNX conventions).
class HammingWindow : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:HammingWindow";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &size, bool periodic = true, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &size, bool periodic, Tensor &output) const;

  /// Output is a float vector while the input is an int scalar: storage
  /// can never be shared with an input.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// General matrix multiplication: Y = alpha * op(A) * op(B) + beta * C.
/// ``transA``/``transB`` control whether A and B are transposed (0 = no,
/// non-zero = yes); ``alpha`` and ``beta`` are scalar multipliers.
/// When ``c`` is ``nullptr`` the bias term is omitted (treated as zero).
class Gemm : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Gemm";
  explicit Gemm(const KernelContext &ctx);
  /// Registers one portable tuning schema for every supported element type.
  static void RegisterTuningSchemas();
  KernelTuningKey TuningKey(int32_t element_type) const override;
  void Configure(const KernelTuningParameters &parameters) override;
  void Run(RuntimeContext &rt) override;
  Tensor operator()(const Tensor &a, const Tensor &b, const Tensor *c, float alpha, float beta,
                    int64_t transA, int64_t transB, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &a, const Tensor &b, const Tensor *c, float alpha, float beta,
                  int64_t transA, int64_t transB, Tensor &output) const;

  /// Returns the immutable configuration used by the execution path.
  const tuning::GemmTuning &tuning() const noexcept { return tuning_; }

  /// Gemm produces a new matrix that cannot alias any input.
  static constexpr bool CanRunInPlace() noexcept { return false; }

private:
  tuning::GemmTuning tuning_;
};

/// Matrix product that behaves like NumPy/ONNX ``matmul``.
///
/// Supports rank-1 and rank-N inputs:
/// - rank-1 x rank-1 -> scalar
/// - rank-2 x rank-2 -> matrix
/// - higher-rank prefixes are broadcast, then batched matrix multiply
class MatMul : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:MatMul";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &a, const Tensor &b, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &a, const Tensor &b, Tensor &output) const;

  /// MatMul generally changes shape and cannot alias inputs safely.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference ``MatMulInteger`` kernel (opset 10). Computes
/// ``Y = matmul(A - a_zero_point, B - b_zero_point)`` with INT8/UINT8 inputs
/// and an INT32 output. ``a_zero_point`` and ``b_zero_point`` are optional and
/// default to 0 when default-constructed (empty) tensors are passed. This
/// reference implementation handles per-tensor (scalar or one-element 1-D)
/// zero points.
class MatMulInteger : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:MatMulInteger";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;

  /// Returning overload allocating an INT32 output. ``a_zero_point`` /
  /// ``b_zero_point`` may be default-constructed (empty) tensors meaning the
  /// optional input is absent.
  Tensor operator()(const Tensor &a, const Tensor &b, const Tensor &a_zero_point,
                    const Tensor &b_zero_point, RuntimeContext *rt = nullptr) const;

  /// In-place overload writing into a caller-allocated INT32 output.
  void operator()(const Tensor &a, const Tensor &b, const Tensor &a_zero_point,
                  const Tensor &b_zero_point, Tensor &output) const;

  /// Output dtype (INT32) differs from inputs (INT8/UINT8), so storage cannot
  /// be shared with an input.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Element-wise sum of a list of tensors with NumPy-style (multidirectional)
/// broadcasting. At least one input is required. All inputs must share the
/// same float dtype (FLOAT or DOUBLE); the output has the broadcast shape of
/// all inputs and the same dtype.
class Sum : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Sum";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensors &inputs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensors &inputs, Tensor &output) const;

  /// Variadic element-wise kernel: the output buffer may alias an input
  /// buffer when that input is not broadcast-expanded (i.e. its shape equals
  /// the output shape).
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise maximum of a list of tensors with NumPy-style
/// (multidirectional) broadcasting. At least one input is required. All
/// inputs must share the same numeric dtype (FLOAT, DOUBLE, INT8, INT16,
/// INT32, INT64, UINT8, UINT16, UINT32 or UINT64); the output has the
/// broadcast shape of all inputs and the same dtype.
class Max : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Max";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensors &inputs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensors &inputs, Tensor &output) const;

  /// Variadic element-wise kernel: the output buffer may alias an input
  /// buffer when that input is not broadcast-expanded.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise minimum of a list of tensors with NumPy-style
/// (multidirectional) broadcasting. At least one input is required. All
/// inputs must share the same numeric dtype (FLOAT, DOUBLE, INT8, INT16,
/// INT32, INT64, UINT8, UINT16, UINT32 or UINT64); the output has the
/// broadcast shape of all inputs and the same dtype.
class Min : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Min";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensors &inputs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensors &inputs, Tensor &output) const;

  /// See :class:`Max`.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Element-wise mean of a list of tensors with NumPy-style (multidirectional)
/// broadcasting. At least one input is required. All inputs must share the
/// same float dtype (FLOAT or DOUBLE); the output has the broadcast shape of
/// all inputs and the same dtype. The result is the element-wise sum of all
/// inputs divided by the number of inputs.
class Mean : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Mean";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensors &inputs, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensors &inputs, Tensor &output) const;

  /// Variadic element-wise kernel: the output buffer may alias an input
  /// buffer when that input is not broadcast-expanded.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// MelWeightMatrix generates the triangular Mel filter-bank weight matrix
/// used to re-weight a linearly sampled frequency spectrum into Mel-scaled
/// bands (since opset 17). The output has shape
/// ``[floor(dft_length/2) + 1, num_mel_bins]`` with element type
/// ``output_dtype`` (default ``FLOAT``).
class MelWeightMatrix : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:MelWeightMatrix";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &num_mel_bins, const Tensor &dft_length, const Tensor &sample_rate,
                    const Tensor &lower_edge_hertz, const Tensor &upper_edge_hertz,
                    DataType output_dtype = DataType::FLOAT, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &num_mel_bins, const Tensor &dft_length, const Tensor &sample_rate,
                  const Tensor &lower_edge_hertz, const Tensor &upper_edge_hertz,
                  DataType output_dtype, Tensor &output) const;

  /// Output shape is unrelated to input shapes (all inputs are scalars):
  /// storage cannot be shared with an input.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};
/// INT32 or INT64 tensor whose value selects the dimension along which the
/// cumulative sum is computed (negative values count from the back). The
/// ``exclusive`` flag, when true, excludes the current element from each
/// position's running sum (the j-th output element is the sum of the first
/// j-1 elements; the 0-th element becomes 0). The ``reverse`` flag, when
/// true, performs the summation in the opposite direction along ``axis``.
class CumSum : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:CumSum";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, const Tensor &axis, bool exclusive = false,
                    bool reverse = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &axis, bool exclusive, bool reverse,
                  Tensor &output) const;

  /// Each output element depends on a previous output element along
  /// ``axis``; aliasing input and output buffers is safe because we
  /// always read the current element before writing the output at
  /// the same position. The ``exclusive`` mode reads the *previous*
  /// position's input before writing, which is also safe in row-major
  /// order when iterating outward from the starting end.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Cumulative product of the input tensor along ``axis``. Semantics mirror
/// :class:`CumSum` with addition replaced by multiplication; in
/// ``exclusive`` mode the starting value is 1 (multiplicative identity).
class CumProd : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:CumProd";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &x, const Tensor &axis, bool exclusive = false,
                    bool reverse = false, RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, const Tensor &axis, bool exclusive, bool reverse,
                  Tensor &output) const;

  /// See :class:`CumSum`.
  static constexpr bool CanRunInPlace() noexcept { return true; }
};

/// Precomputed contraction plan for :cpp:class:`Einsum` (defined in the
/// implementation file). Cached between calls so it is built only once per
/// distinct equation/input-shape combination.
struct EinsumPlan;

/// Reference implementation of the ``Einsum`` operator (opset 12).
///
/// Evaluates the Einstein summation expressed by ``equation`` over the list
/// of input tensors. The equation may contain ellipsis (``...``) to broadcast
/// leading dimensions, and may be given either in explicit form (``->``
/// followed by the output term) or implicit form. All inputs must share the
/// same dtype (FLOAT or DOUBLE); the output has the same dtype.
class Einsum : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:Einsum";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensors &inputs, const std::string &equation,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensors &inputs, const std::string &equation, Tensor &output) const;

  /// Einsum generally changes shape and cannot alias inputs safely.
  static constexpr bool CanRunInPlace() noexcept { return false; }

private:
  /// Returns the contraction plan for ``inputs``/``equation``, building it once
  /// and reusing the cached result while the equation and input shapes are
  /// unchanged, so :cpp:func:`BuildPlan` is not recomputed on every call.
  const EinsumPlan &EnsurePlan(const Tensors &inputs, const std::string &equation) const;

  mutable std::shared_ptr<const EinsumPlan> plan_;
  mutable std::string plan_equation_;
  mutable std::vector<Shape> plan_shapes_;
};

/// Reference implementation of the ONNX ``DFT`` operator (since opset 17;
/// in opset 20 the ``axis`` attribute moved to an optional third input).
///
/// Computes the (one-dimensional) discrete Fourier Transform of ``input``
/// along ``axis``. ``input`` must be a floating-point tensor of rank >= 2,
/// where the trailing dimension is ``1`` (real-valued samples) or ``2``
/// (interleaved real/imaginary parts). The returned tensor has the same
/// rank as ``input``; its last dimension is ``2`` (complex output) except
/// for IRFFT (``onesided=1`` and ``inverse=1``), in which case it is ``1``.
///
/// When ``dft_length`` is specified the signal is zero-padded or truncated
/// along ``axis``; otherwise the axis dimension is used (or
/// ``2 * (signal_dim_axis - 1)`` for the IRFFT default).
class DFT : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:DFT";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  /// ``axis`` is the signal axis (must satisfy ``-rank <= axis``, ``axis !=
  /// -1`` and ``axis < rank - 1``). ``dft_length`` is a pointer to a 0-D
  /// INT32/INT64 tensor; pass ``nullptr`` to use the default.
  Tensor operator()(const Tensor &input, const Tensor *dft_length, int64_t axis,
                    bool onesided = false, bool inverse = false,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &input, const Tensor *dft_length, int64_t axis, bool onesided,
                  bool inverse, Tensor &output) const;

  /// The output shape (last dim and signal axis) generally differs from the
  /// input shape, so storage cannot be shared.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

/// Reference implementation of the ONNX ``STFT`` operator (opset 17).
///
/// Computes the Short-time Fourier Transform of ``signal`` by sliding a
/// (optionally windowed) DFT of length ``frame_length`` over the signal axis
/// with step ``frame_step``. ``signal`` must be a floating-point tensor of
/// shape ``[batch_size, signal_length, 1]`` (real input) or
/// ``[batch_size, signal_length, 2]`` (interleaved real/imaginary parts).
/// The output has shape ``[batch_size, n_frames, dft_unique_bins, 2]`` where
/// ``n_frames = (signal_length - frame_length) / frame_step + 1`` and
/// ``dft_unique_bins`` is ``floor(frame_length / 2) + 1`` when ``onesided``
/// is true (the default for real input) or ``frame_length`` otherwise.
///
/// ``window`` is an optional 1-D tensor with shape ``[frame_length]``; pass
/// ``nullptr`` to skip the windowing step. ``frame_length`` is an optional
/// pointer to a 0-D INT32/INT64 tensor; pass ``nullptr`` to derive
/// ``frame_length`` from the ``window`` shape. At least one of ``window``
/// or ``frame_length`` must be provided.
class STFT : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:STFT";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  Tensor operator()(const Tensor &signal, const Tensor &frame_step, const Tensor *window,
                    const Tensor *frame_length, bool onesided = true,
                    RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &signal, const Tensor &frame_step, const Tensor *window,
                  const Tensor *frame_length, bool onesided, Tensor &output) const;

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

/// Reference implementation of the ``TopK`` operator (opsets 1, 10, 11).
///
/// Selects the ``k`` largest (or smallest, when ``largest`` is false) values
/// along ``axis`` from the input tensor ``x`` and returns them together with
/// the corresponding indices. ``axis`` is interpreted modulo the rank of
/// ``x`` (negative values count from the back). When ``sorted`` is true (the
/// default) the returned values are sorted descending (or ascending when
/// ``largest`` is false); ties along ``axis`` are broken by the lower index.
/// When ``sorted`` is false the order of the returned values is unspecified
/// by the ONNX schema; this reference implementation still returns them
/// sorted to keep the output deterministic.
///
/// Returns a ``std::pair<Tensor, Tensor>`` where the first tensor (``Values``)
/// has the same dtype as ``x`` and the second tensor (``Indices``) is an
/// ``INT64`` tensor. Both share the shape of ``x`` with the ``axis``
/// dimension replaced by ``k``.
class TopK : public KernelBase {
public:
  static constexpr const char *name = "onnx_kernels:CPU:ai.onnx:TopK";
  void Run(RuntimeContext &rt) override;
  using KernelBase::KernelBase;
  std::pair<Tensor, Tensor> operator()(const Tensor &x, int64_t k, int64_t axis = -1,
                                       bool largest = true, bool sorted = true,
                                       RuntimeContext *rt = nullptr) const;
  void operator()(const Tensor &x, int64_t k, int64_t axis, bool largest, bool sorted,
                  Tensor &values, Tensor &indices) const;

  /// TopK output shape differs from the input along ``axis`` (k vs. axis_dim)
  /// and the Indices output has a different dtype (int64), so the output
  /// buffers cannot safely alias either input buffer.
  static constexpr bool CanRunInPlace() noexcept { return false; }
};

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