include_controlflow_kernels.h#

namespace onnx_light

Alias that makes onnx-light headers compatible with code that references ONNX_LIGHT_NAMESPACE (the macro used in the standard onnx package).

Set to ONNX_LIGHT_NAMESPACE so both names resolve to the same namespace.

Symbol-visibility attribute for the public onnx-light C++ API.

Defined as empty because onnx-light does not require explicit __declspec(dllexport) or __attribute__((visibility("default"))) annotations — visibility is controlled at the shared-library level. The macro is provided so that vendored ONNX headers that decorate their declarations with ONNX_API compile without modification.

Namespace alias so that ONNX C++ code (and consumers such as onnxruntime) that refers to the literal onnx namespace — rather than the ONNX_NAMESPACE macro — resolves to the onnx-light namespace. The standard onnx package lives in namespace onnx; onnx-light uses onnx_light (via ONNX_LIGHT_NAMESPACE), so this alias keeps onnx-light a true drop-in. It is only introduced when the onnx-light namespace differs from onnx.

namespace onnx_kernels
namespace kernel
class If : public onnx_light::onnx_kernels::kernel::KernelBase#
#include <include_controlflow_kernels.h>

Selects then_value when the scalar BOOL cond is true, otherwise returns else_value. Both branch values must share the same data type and shape.

Public Functions

Tensor operator()(const Tensor &cond, const Tensor &then_value, const Tensor &else_value) const#
void operator()(const Tensor &cond, const Tensor &then_value, const Tensor &else_value, Tensor &output) const#
std::vector<Tensor> operator()(const Tensor &cond, const GraphProto &then_branch, const GraphProto &else_branch, RuntimeContext &rt) const#

Branch-graph overload.

Selects the then_branch :cpp:class:GraphProto when the scalar BOOL cond is true and else_branch otherwise, then executes the selected subgraph through :cpp:func:RunGraph using rt (a child :cpp:class:RuntimeContext is created internally so the caller’s tensor map is left untouched apart from being inherited as the outer scope of the subgraph). The returned vector contains the subgraph outputs in the order declared by branch.output().

Throws:

std::invalid_argument – if cond is not a BOOL scalar, if a declared subgraph output is missing, if then_branch and else_branch declare a different number of outputs, or if any node of the executed subgraph fails to dispatch.

inline explicit KernelBase(const KernelContext &ctx)#

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept#
class Loop : public onnx_light::onnx_kernels::kernel::KernelBase#
#include <include_controlflow_kernels.h>

Reference implementation of the ONNX Loop operator.

Like :class:If, this kernel does not execute the loop body subgraph itself; it consumes already-evaluated per-iteration values and merely validates and assembles the operator’s outputs:

  • each of the N final loop-carried dependency values is forwarded verbatim from the caller-provided final_state tensors;

  • each of the K scan outputs is built by stacking the caller-provided per-iteration values along a new leading axis whose length equals the actually executed trip count.

The kernel honors ONNX’s termination rules: the effective trip count is min(M, len(scan_values_per_iter)) when M is provided, otherwise len(scan_values_per_iter); when cond is false on entry, the trip count is zero and scan outputs are zero-length along the new axis. The kernel is therefore a faithful reference for the operator’s composition/stacking semantics that is useful for shape and type-propagation tests while keeping the implementation independent from any graph executor.

Public Types

using BodyRunner = std::function<std::vector<Tensor>(int64_t iter, bool cond_in, const std::vector<Tensor> &state)>#

Body-runner callback executed by the new operator() overload once per iteration. iter is the 0-based iteration index, cond_in is the BOOL termination condition entering this iteration, and state is the current loop-carried state (v_initial for iter == 0, otherwise the previous iteration’s body outputs [1..1+N)). The callback must return 1 + N + K tensors matching the ONNX Loop body output convention: the new cond_out (BOOL scalar) at index 0, the N updated loop-carried values at indices [1, 1+N), and the K per-iteration scan values at indices [1+N, 1+N+K).

Public Functions

std::vector<Tensor> operator()(const Tensor &M, const Tensor &cond, const std::vector<Tensor> &v_initial, const std::vector<Tensor> &final_state, const std::vector<std::vector<Tensor>> &scan_values_per_iter) const#

Returning overload.

Parameters:

  • M – Optional INT64 scalar maximum trip-count (empty Tensor means omitted).

  • cond – Optional BOOL scalar initial termination condition (empty Tensor means omitted, treated as true).

  • v_initial – Initial loop-carried dependency values (size N, may be empty). Returned as-is when the loop executes zero iterations.

  • final_state – Final loop-carried dependency values (size N). Their data types must match the corresponding v_initial tensor’s data type.

  • scan_values_per_iter – Per-iteration scan-output values, given as scan_values_per_iter[k][t] &#8212; K scan outputs each with one tensor per iteration (rectangular). All entries within a scan-output row must share the same data type and shape.

Returns:

N + K tensors: the final loop-carried dependency values followed by the stacked scan outputs.

std::vector<Tensor> operator()(const Tensor &M, const Tensor &cond, const std::vector<Tensor> &v_initial, std::size_t num_scan_outputs, const BodyRunner &run_body) const#

Body-aware overload.

Performs the full ONNX Loop semantics: honors M / cond termination rules, invokes run_body once per iteration, threads the loop-carried state across iterations, collects each iteration’s scan values, and assembles the N + num_scan_outputs outputs (final loop-carried state followed by stacked scan outputs).

Parameters:

  • M – Optional INT64 scalar maximum trip-count (empty Tensor means omitted).

  • cond – Optional BOOL scalar initial termination condition (empty Tensor means omitted, treated as true).

  • v_initial – Initial loop-carried dependency values (size N).

  • num_scan_outputs – Number of per-iteration scan outputs the body produces (K).

  • run_body – Callback executed once per iteration; see :type:BodyRunner.

Returns:

N + num_scan_outputs tensors.

inline explicit KernelBase(const KernelContext &ctx)#

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept#
class Scan : public onnx_light::onnx_kernels::kernel::KernelBase#
#include <include_controlflow_kernels.h>

Reference implementation of the ONNX Scan operator.

Two operator() overloads are provided:

  • a body-aware overload that takes the Scan body subgraph, the initial state and the per-axis scan inputs, executes the body once per iteration through :cpp:func:onnx_kernels::RunSubgraph, and returns the operator’s full output list. This is the overload used when the kernel is invoked from the runtime dispatcher (:cpp:func:onnx_kernels::RunNode);

  • a stacking-only overload that consumes already-evaluated state and per-iteration values and merely assembles the operator’s outputs. It is retained for legacy callers and for tests that want to exercise the stacking semantics in isolation, without standing up a runtime context.

For both overloads:

  • each of the N final state values is returned as either the final body-produced state (when the trip count is non-zero) or the corresponding initial_state entry (when the trip count is zero);

  • each of the K scan outputs is built by stacking the per-iteration values along a new axis whose position is scan_output_axes[k] (default 0 &#8212; new leading axis) and whose length equals the trip count. When scan_output_directions[k] equals 1 the per-iteration values are reversed before stacking (prepend semantics).

Public Functions

std::vector<Tensor> operator()(const GraphProto &body, const std::vector<Tensor> &initial_state, const std::vector<Tensor> &scan_inputs, RuntimeContext &rt, const std::vector<int64_t> &scan_input_axes = {}, const std::vector<int64_t> &scan_input_directions = {}, const std::vector<int64_t> &scan_output_axes = {}, const std::vector<int64_t> &scan_output_directions = {}) const#

Body-aware overload.

Iterates the Scan body subgraph for the trip count implied by the scan inputs, threading the state forward across iterations and collecting the per-iteration scan outputs, then returns the stacked operator outputs.

Parameters:

  • body – The Scan body subgraph. Its first N formal inputs are bound to the current state, its next M formal inputs are bound to the per-iteration scan-input slices, its first N outputs become the next state, and its remaining K outputs are the per-iteration scan outputs.

  • initial_state – Initial state values (size N).

  • scan_inputs – Per-axis scan inputs (size M). Each must have rank at least 1; the axis given by scan_input_axes is consumed one element at a time and determines the trip count.

  • rt – Runtime context used to evaluate body. The body is executed in a fresh child context per iteration so it cannot mutate rt.tensors().

  • scan_input_axes – Per-scan-input axis along which the input is sliced. When empty, axis 0 is used for every scan input. Negative values count from the back of the scan input’s rank. When non-empty, must have M entries.

  • scan_input_directions – Per-scan-input direction (0 = forward, 1 = reverse). When empty, all scan inputs use the forward direction. When non-empty, must have M entries.

  • scan_output_axes – Per-scan-output axis at which the new stacking axis is inserted (see the stacking-only overload).

  • scan_output_directions – Per-scan-output direction (see the stacking-only overload).

Returns:

N + K tensors: the final state values followed by the stacked scan outputs.

std::vector<Tensor> operator()(int64_t trip_count, const std::vector<Tensor> &initial_state, const std::vector<Tensor> &final_state, const std::vector<std::vector<Tensor>> &scan_values_per_iter, const std::vector<int64_t> &scan_output_axes = {}, const std::vector<int64_t> &scan_output_directions = {}) const#

Stacking-only overload.

Parameters:

  • trip_count – Number of iterations actually executed (non-negative).

  • initial_state – Initial state values (size N).

  • final_state – Final state values (size N). Their data types must match the corresponding initial_state entry’s data type.

  • scan_values_per_iter – Per-iteration scan-output values, given as scan_values_per_iter[k][t] — K scan outputs each with one tensor per iteration (rectangular). All entries within a scan-output row must share the same data type and shape.

  • scan_output_axes – Per-scan-output axis at which the new stacking axis is inserted. When empty, axis 0 is used for every scan output. Negative values count from the back of the stacked output (with rank elt.rank + 1). When non-empty, must have K entries.

  • scan_output_directions – Per-scan-output direction (0 = append, 1 = prepend / reverse before stacking). When empty, all scan outputs use the append direction. When non-empty, must have K entries.

Returns:

N + K tensors: the final state values followed by the stacked scan outputs.

inline explicit KernelBase(const KernelContext &ctx)#

Public Static Functions

static inline constexpr bool CanRunInPlace() noexcept#