torch2casadi

Export a PyTorch model and its derivative graphs to ONNX, then use it in CasADi optimization. PyTorch is needed at export time; evaluation uses CasADi's native ONNX Runtime backend.

An MIT-licensed early-development package.

Source: https://github.com/casadi/torch2casadi

import torch
import casadi as ca
from torch2casadi import export

model = torch.nn.Sequential(
    torch.nn.Linear(2, 16), torch.nn.Tanh(), torch.nn.Linear(16, 2)
).eval()
path = export(model, torch.zeros(1, 2), "generated")
f = ca.GraphBuilder(str(path)).create("f")

opti = ca.Opti()
u = opti.variable(2)
opti.minimize(ca.sumsqr(f(u) - ca.DM([0.9, -0.2])))
opti.subject_to(opti.bounded(-1, u, 1))
opti.solver("sqpmethod", {"qpsol": "qrqp"})
solution = opti.solve()

The trained-surrogate example in examples/opti.py includes a coupled constraint and regularization. It uses exact Hessians.

Installation

Use Python 3.10–3.12:

python -m pip install torch2casadi

Dependencies are deliberately user-managed: installing torch2casadi does not install or upgrade Torch, CasADi, or any other runtime package. Export needs Torch, ONNX, ONNXScript and NumPy in your environment. The Python ONNX Runtime package is needed only for numerical tests. Keep the Torch build appropriate for your hardware and workflow.

CI currently tests Torch 2.6 with the companion versions listed in requirements-ci.txt. That file describes a tested environment, not installation requirements or a guarantee about other versions. The exporter uses experimental Torch APIs, so custom builds need validation.

For development, install the checkout with python -m pip install -e . and install build and twine separately if needed.

The published 0.1 release predates this policy and still declares runtime dependencies. To install that version without changing an existing environment, use python -m pip install --no-deps torch2casadi==0.1. Starting with 0.1.1, no runtime dependencies are declared.

For CasADi integration, build branch onnx-primal-efficiency with WITH_ONNX=ON and WITH_ONNX_RUNTIME=ON. The sibling derivative feature is in commit b5040ffd76. Installing the Python onnxruntime wheel does not enable CasADi's native backend. The exporter itself does not depend on the CasADi Python package.

Export contract

• One tensor input, or a tuple of positional tensor inputs, and one tensor output; float32 or float64 on CPU during export. All inputs use the same dtype.
• Tensor shapes are fixed by the examples. Inputs are differentiable by default, including all model batch elements. is_diff_in can disable differentiation for whole input arguments. No independence between batch elements is assumed.
• The module is copied, switched to evaluation mode and frozen. The caller's parameters and training mode are preserved.
• Inputs and outputs are flattened in PyTorch element order and exposed as CasADi column vectors. Derivative direction counts are dynamic.
• Modules must support the required PyTorch function transforms and ONNX operations. Data-dependent Python branches, stochastic evaluation, integer/discrete inputs, nested argument/output trees, and unsupported custom operators are outside this initial contract.
• Smooth models are appropriate for exact-Hessian optimization. ReLU networks do not acquire smooth Hessians through export.

Default files:

export(..., forward=True) also writes fwd_f.onnx. export(..., hessian=False) skips fwd_adj_f.onnx. name="net" changes the family to net.onnx, adj_net.onnx, etc.

Combined columns are grouped by outer forward direction, then inner adjoint direction. The mixed graph differentiates both the primal input and the adjoint seed. It computes products through AD without forming a dense Jacobian or Hessian as an intermediate.

Each export checks all graphs with ONNX's checker. Files are staged until all graphs pass. Export into a fresh model-family directory, or pass overwrite=True to replace an existing family. Derivative files no longer produced are removed; unrelated models are preserved.

ONNX Runtime is a test dependency only; it is not imported or required by the exporter. The tests execute exported graphs at perturbed inputs with independent (nadj, nfwd) counts (1,1), (2,3), (3,2), (4,5) and compare primal, adjoint, forward, and mixed derivatives with PyTorch AD. Tolerances are 2e-5/2e-6 for float32 and 1e-8/1e-9 for float64. These checks cover the test models, not every model supplied to export().

Keep sibling files together until derivatives are constructed. Constructed CasADi derivatives embed their model bytes, including when serialized.

Runtime parameters without derivatives

Expose a prescribed quantity, such as ambient temperature, as a separate model argument:

class Vibration(torch.nn.Module):
    def __init__(self, network):
        super().__init__()
        self.network = network  # three features -> one vibration score

    def forward(self, x, temperature):
        return self.network(torch.cat((x, temperature), dim=1))

flags = [True, False]
path = export(Vibration(network), (torch.zeros(1, 2), torch.zeros(1, 1)), "generated",
              input_names=["x", "temperature"], is_diff_in=flags)
r = ca.GraphBuilder(str(path)).create("r")

opti = ca.Opti()
x = opti.variable(2)
temperature = opti.parameter()
opti.set_value(temperature, 0.0)
opti.subject_to(r(x, temperature) <= 0.8)

With the companion CasADi inference patch, numeric import infers the mask from the forward seed inputs or adjoint sensitivity outputs in the derivative family. Explicit is_diff_in overrides inference; conflicting forward/adjoint masks are rejected. Without derivative signatures, inputs remain differentiable by default. Inference reads sibling signatures during construction; supplying a mask keeps sibling loading lazy. For symbolic import, use {"symbolic": True, "is_diff_in": flags} with a CasADi build containing the GraphBuilder option forwarding fix.

The primal and derivative graphs still accept temperature, so changing its value changes both predictions and state derivatives. The exporter creates neither adj_temperature outputs nor fwd_temperature seeds. Forward-over-adjoint differentiates only the active state inputs and the adjoint seed. CasADi's full derivative calling convention still includes the omitted slots; its wrapper supplies zeros for nondifferentiable sensitivities. This intentionally does not provide temperature or optimal-solution sensitivities.

For this example the files are:

forward=True additionally exports fwd_f.onnx with inputs x, temperature, fwd_x and output fwd_y. Multiple differentiable arguments produce separate adjoint outputs and forward seeds. A single input defaults to the name x; multiple inputs default to x0, x1, etc. An all-false mask exports only the primal model.

examples/shuttle.py trains a toy temperature-dependent vibration model and solves a four-stage multiple-shooting problem with FATROP, interleaved state/control variables and automatic structure detection. It checks both numeric and symbolic ONNX paths at two temperature values. Its model and data are illustrative, not a calibrated physical model.

Weights as decision variables

By default the model parameters are baked into the graphs as constants, which is what inference needs. To train inside a CasADi NLP (a universal ODE, a surrogate with hard constraints, ...), the weights must stay decision variables. No exporter support is needed for that: make the weights arguments of forward, and they become graph inputs of the whole family like any other argument. A small wrapper does this for an existing module:

class FreeWeights(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.names = [n for n, _ in model.named_parameters()]

    def forward(self, x, *params):
        return torch.func.functional_call(self.model, dict(zip(self.names, params)), (x,))

model = torch.nn.Sequential(torch.nn.Linear(1, 8), torch.nn.Tanh(), torch.nn.Linear(8, 1))
params = [p.detach() for p in model.parameters()]
path = export(FreeWeights(model), (torch.zeros(1, 1), *params), "generated",
              input_names=["v", "W1", "b1", "W2", "b2"], hessian=False)
f = ca.GraphBuilder(str(path)).create("f")   # f:(v,W1[8],b1[8],W2[8],b2)->(y)

opti = ca.Opti()
weights = [opti.variable(p.numel()) for p in params]
for w, p in zip(weights, params):
    opti.set_initial(w, p.reshape(-1).numpy())
...
opti.solver("ipopt", {}, {"hessian_approximation": "limited-memory",
                          "limited_memory_max_history": 50})
sol = opti.solve()
for w, p in zip(weights, model.parameters()):
    p.data = torch.tensor(sol.value(w)).reshape(p.shape).to(p.dtype)

Pass a subset of named_parameters() to free only some layers. Each weight tensor is one input, flattened in PyTorch element order like every other input. Gradients with respect to the weights come from adj_f.onnx at the cost of one adjoint sweep, whatever their number. Exact Hessians are different: fwd_adj_f.onnx is swept once per forward direction, and a Hessian block over the weights needs as many directions as there are weights. For small networks that is fine and gives Newton-type convergence; for large ones pass hessian=False and use a quasi-Newton Hessian in the NLP solver, as above. The primal graph also imports symbolically ({"symbolic": True}), in which case CasADi differentiates it itself.

Development

Run the exporter tests and build the distributions:

python -m unittest discover -s tests
python -m build

For the additional integration tests, put the feature CasADi build on PYTHONPATH and run TORCH2CASADI_INTEGRATION=1 python -m unittest discover -s tests. Those tests require the native ORT plugin and are explicitly skipped in ordinary package CI until the CasADi feature is released.

The tests cover Tanh, Sigmoid, Softplus, tensor-shaped nonlinear algebra, independent seed counts, Hessians through CasADi and serialization after the ONNX files are deleted.

The implementation uses experimental FX tracing APIs and a PyTorch 2.6 export workaround. Broader PyTorch-version/operator support needs a compatibility test matrix before a stable release. Future arbitrary derivative-family export can build on the same packaging convention; this initial package emits first derivatives and forward-over-adjoint, not all nesting patterns supported by CasADi.

Releases and provenance

The release guide describes TestPyPI/PyPI Trusted Publishing and verification. Releases use OIDC with GitHub build-provenance attestations and PyPI publishing attestations.