onnx

import "github.com/umbralcalc/stochadex/pkg/onnx"

Package onnx is an opt-in inference partition: it runs a frozen ONNX model behind the engine’s simulator.Iteration interface, so a model trained upstream (in Python — sklearn, XGBoost, a small neural net — and exported to .onnx) can be a component in a stochadex simulation. Each step the partition reads a feature vector from its params, runs the model through a cgo ONNX Runtime session, and returns the model output as its next state.

Why a separate module

This is a SEPARATE module on purpose, exactly like pkg/duckdbstore and pkg/s3store: it pulls in a cgo dependency (github.com/yalue/onnxruntime_go) and the ONNX Runtime shared library, neither of which the engine’s own go.mod carries. The engine core therefore stays lean and CGO_ENABLED=0-clean (it cross-compiles and builds to WASM for everyone who imports it as a library), while consumers who want inference opt in by importing this module. The implementation is behind the `onnx` build tag and CGO must be enabled.

How the opt\-in works

Importing this module (under the onnx tag) self-registers the {type: onnx_inference} spelling with the engine’s config surface through api.RegisterIteration — the same downstream-registration hook the Arrow, S3 and DuckDB spellings use (RegisterDataSource / simulator.RegisterComponent). The engine core never imports this package; a CLI or a downstream library reaches it with a blank import, e.g. cmd/stochadex does so under its onnx tag.

Config surface

A single-input model uses the shorthand:

iteration:
  type: onnx_inference
  model_path: model.onnx        # required
  input_param: input            # params key holding the feature vector (default "input")
  input_name: ...               # ONNX graph input name (default: the sole input)
  output_name: ...              # ONNX graph output name (default: the sole output)
  shared_library_path: ...      # ONNX Runtime library (see below)
  intra_op_threads: 1           # ONNX Runtime intra-op pool size (see Threading)
  inter_op_threads: 1           # ONNX Runtime inter-op pool size

A multi-input model binds each model input to a params key with an inputs map (mutually exclusive with input_param / input_name):

iteration:
  type: onnx_inference
  model_path: model.onnx
  inputs:                       # {params key: ONNX input name}, every input bound
    features: input
    theta: parameters

Each input vector is wired in like any other partition input — a static params: entry, or params_from_upstream from a producing partition. The model output becomes the partition’s state, so its state_width must equal the flattened output length. Inputs and the output are float32 tensors (the common export dtype); dynamic dimensions such as a symbolic batch size are pinned to 1 for single-step inference.

Tuning model parameters

Because inputs are just params vectors, a model exported with some inputs designated as *parameters* rather than features (a hand-built graph or a torch export — skl2onnx will not do this) becomes tunable by the framework’s optimisation / SBI tools with no engine change: bind the parameter input to its own params key and wire that key from a partition the sampler perturbs. The model stays frozen; only the parameter vector arriving through params moves.

Threading

intra_op_threads / inter_op_threads size the session’s ONNX Runtime thread pools; unset (<= 0) leaves the runtime default. For single-row inference across many concurrent partitions, each session’s default pool oversubscribes the CPU and the intra-op parallelism buys nothing — setting both to 1 is usually the better choice. The knob is exposed so the caller tunes it per partition rather than the engine guessing a policy.

ONNX Runtime shared library

The library is a system dependency (mirroring the cblas “CGO with a system BLAS” story rather than vendoring per-platform binaries), located at run time in this order: the spec’s shared_library_path, then the ONNXRUNTIME_LIB_PATH environment variable, then a short list of conventional install locations, then the binding’s own default search.

Scope — inference only

The partition is a pure, allocation-free map from input to prediction: all tensors are allocated once in Configure and reused every step (copy into a pinned input slice, run, copy out). Fitting and refitting stay upstream, per the engine’s generative/inferential repo boundary — a downstream repo hands the engine a new .onnx artifact; this package never trains.

Index

• func BuildIteration(spec simulator.ComponentSpec) (simulator.Iteration, error)
• type OnnxInferenceIteration
  • func (o *OnnxInferenceIteration) Configure(partitionIndex int, settings *simulator.Settings)
  • func (o *OnnxInferenceIteration) Iterate(params *simulator.Params, partitionIndex int, stateHistories []*simulator.StateHistory, timestepsHistory *simulator.CumulativeTimestepsHistory) []float64
• type OnnxInput

func BuildIteration

func BuildIteration(spec simulator.ComponentSpec) (simulator.Iteration, error)

BuildIteration constructs an OnnxInferenceIteration from a data spec, validating fields strictly (an unknown key is an error, matching the rest of the config surface). Inputs are given either as the single-input shorthand (input_param / input_name) or as a multi-input map (inputs: {param: model_input}); the two forms are mutually exclusive.

type OnnxInferenceIteration

OnnxInferenceIteration runs a frozen ONNX model as a partition. Each step it reads each bound input vector from params, runs the model, and returns the flattened model output as the next state. All tensors are allocated once in Configure and reused every step (Invariant B): Iterate copies into pinned []float32 inputs, runs, and copies the output back out — no per-step allocation.

type OnnxInferenceIteration struct {
    // Config (set by the builder, immutable across runs).
    ModelPath         string
    OutputName        string
    SharedLibraryPath string
    // Inputs binds params keys to named model inputs. When empty, the single-input
    // shorthand below is used instead. Multi-input unlocks parameter tuning: a
    // parameter vector arrives as its own input alongside the features.
    Inputs []OnnxInput
    // InputParam / InputName are the single-input shorthand, used only when Inputs
    // is empty. InputParam defaults to "input"; InputName defaults to the sole
    // model input.
    InputParam string
    InputName  string
    // IntraOpThreads / InterOpThreads size the session's ONNX Runtime thread pools.
    // <= 0 leaves the runtime default. For single-row inference across many
    // concurrent partitions, setting both to 1 avoids thread oversubscription (each
    // session otherwise spins up its own default pool); the caller tunes this
    // rather than the engine guessing.
    IntraOpThreads int
    InterOpThreads int
    // contains filtered or unexported fields
}

func (*OnnxInferenceIteration) Configure

func (o *OnnxInferenceIteration) Configure(partitionIndex int, settings *simulator.Settings)

func (*OnnxInferenceIteration) Iterate

func (o *OnnxInferenceIteration) Iterate(params *simulator.Params, partitionIndex int, stateHistories []*simulator.StateHistory, timestepsHistory *simulator.CumulativeTimestepsHistory) []float64

type OnnxInput

OnnxInput binds one params key to one named model input, so a partition can feed a multi-input model — e.g. a feature vector under one key and a tunable parameter vector under another, the latter driven by the framework’s optimisation / SBI tools exactly like any other partition parameter.

type OnnxInput struct {
    // ParamKey is the params key the input vector is read from each step.
    ParamKey string
    // ModelInputName is the ONNX graph input to feed. Empty means the model's
    // sole input (only valid when the model has exactly one).
    ModelInputName string
}

Generated by gomarkdoc