onnx_function.cpp

/*
 *    This file is part of CasADi.
 *
 *    CasADi -- A symbolic framework for dynamic optimization.
 *    Copyright (C) 2010-2023 Joel Andersson, Joris Gillis, Moritz Diehl,
 *                            KU Leuven. All rights reserved.
 *    Copyright (C) 2011-2014 Greg Horn
 *
 *    CasADi is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU Lesser General Public
 *    License as published by the Free Software Foundation; either
 *    version 3 of the License, or (at your option) any later version.
 *
 *    CasADi is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *    Lesser General Public License for more details.
 *
 *    You should have received a copy of the GNU Lesser General Public
 *    License along with CasADi; if not, write to the Free Software
 *    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */
 
#include "onnx_function_impl.hpp"
#include "graph_builder_internal.hpp"
#include "casadi_misc.hpp"
#include "filesystem_impl.hpp"
 
namespace casadi {
 
  Dict OnnxFunction::info() const {
    Dict ret = FunctionInternal::info();
    ret["model_path"] = model_path_;
    return ret;
  }
 
  bool has_onnxbackend(const std::string& solver) {
    return OnnxFunction::has_plugin(solver);
  }
 
  void load_onnxbackend(const std::string& solver) {
    OnnxFunction::load_plugin(solver);
  }
 
  std::vector<std::string> onnxbackend_solvers() {
    std::vector<std::string> ret;
    for (auto&& s : OnnxFunction::solvers_) ret.push_back(s.first);
    return ret;
  }
 
  std::string onnxbackend_doc(const std::string& solver) {
    return OnnxFunction::getPlugin(solver).doc;
  }
 
  std::map<std::string, OnnxFunction::Plugin> OnnxFunction::solvers_;
 
#ifdef CASADI_WITH_THREADSAFE_SYMBOLICS
  std::mutex OnnxFunction::mutex_solvers_;
#endif // CASADI_WITH_THREADSAFE_SYMBOLICS
 
  const std::string OnnxFunction::infix_ = "onnx";
 
  std::string OnnxFunction::meta_doc = "";
 
  const Options OnnxFunction::options_
  = {{&FunctionInternal::options_},
     {{"provider",
       {OT_STRING, "Execution provider for the ONNX runtime backend"}},
      {"dim_bindings",
       {OT_DICT, "Sizes for symbolic/dynamic tensor dimensions (name -> size)"}},
      {"input_shapes",
       {OT_DICT, "Explicit shapes for inputs (name -> shape), overriding the model's"}},
      {"input_values",
       {OT_DICT, "Baked-in input values (name -> value); these inputs are not exposed"}},
      {"fwd_dim",
       {OT_STRING, "Symbolic dimension naming the forward seed count [nfwd]"}},
      {"adj_dim",
       {OT_STRING, "Symbolic dimension naming the adjoint seed count [nadj]"}}
     }
    };
 
  std::string onnx_dtype_name(casadi_int t) {
    switch (t) {
      case 1: return "FLOAT";       case 2: return "UINT8";    case 3: return "INT8";
      case 4: return "UINT16";      case 5: return "INT16";    case 6: return "INT32";
      case 7: return "INT64";       case 8: return "STRING";   case 9: return "BOOL";
      case 10: return "FLOAT16";    case 11: return "DOUBLE";  case 12: return "UINT32";
      case 13: return "UINT64";     case 16: return "BFLOAT16";
      default: return "TYPE" + str(t);
    }
  }
 
  casadi_int onnx_dtype_enum(const std::string& name) {
    static const std::map<std::string, casadi_int> m = {
      {"FLOAT", 1},  {"UINT8", 2},  {"INT8", 3},   {"UINT16", 4},   {"INT16", 5},
      {"INT32", 6},  {"INT64", 7},  {"STRING", 8}, {"BOOL", 9},     {"FLOAT16", 10},
      {"DOUBLE", 11}, {"UINT32", 12}, {"UINT64", 13}, {"BFLOAT16", 16}};
    auto it = m.find(name);
    return it != m.end() ? it->second : 0;
  }
 
  // Select a subset of tensors by name, preserving the requested order (empty request => all)
  static std::vector<OnnxTensorInfo> select_tensors(const std::vector<OnnxTensorInfo>& all,
                                                    const std::vector<std::string>& req) {
    if (req.empty()) return all;
    std::vector<OnnxTensorInfo> sel;
    for (const std::string& name : req) {
      bool found = false;
      for (const OnnxTensorInfo& t : all)
        if (t.name == name) { sel.push_back(t); found = true; break; }
      casadi_assert(found, "ONNX tensor '" + name + "' not found in model");
    }
    return sel;
  }
 
  OnnxFunction::OnnxFunction(const std::string& name, const GraphBuilderInternal* gb,
                             const std::vector<std::string>& inputs,
                             const std::vector<std::string>& outputs)
    : FunctionInternal(name) {
    // Freeze a snapshot: the builder may mutate afterwards without affecting this function.
    // Shapes are resolved by the builder (dim bindings + input-shape overrides applied).
    model_data_ = gb->model_data_;
    input_values_ = gb->input_values_;
    dim_bindings_ = gb->dim_bindings_;
    input_shapes_ = gb->input_shapes_;
    model_path_ = gb->model_path_;
    builder_opts_ = gb->opts_;
    std::vector<OnnxTensorInfo> all_out;
    for (const Node& n : gb->node_list()) {
      OnnxTensorInfo t;
      t.name = n.name;
      t.elem_type = onnx_dtype_enum(n.dtype);
      t.shape = gb->resolved_shape(n);
      t.numel = 1;
      for (casadi_int d : t.shape) t.numel *= d;
      if (n.io == "input") {
        all_in_.push_back(t);
        model_inputs_.insert(t.name);
      } else {
        all_out.push_back(t);
        model_outputs_.insert(t.name);
      }
    }
    in_ = select_tensors(all_in_, inputs);    // exposed CasADi inputs (selection)
    out_ = select_tensors(all_out, outputs);  // exposed CasADi outputs (selection)
  }
 
  OnnxFunction::~OnnxFunction() {
    clear_mem();
  }
 
  // Pack/unpack a vector<OnnxTensorInfo> as parallel arrays
  static void pack_tensors(SerializingStream& s, const std::string& d,
                           const std::vector<OnnxTensorInfo>& v) {
    std::vector<std::string> names;
    std::vector<std::vector<casadi_int>> shapes;
    std::vector<casadi_int> elem_types, numels;
    for (const OnnxTensorInfo& t : v) {
      names.push_back(t.name);
      shapes.push_back(t.shape);
      elem_types.push_back(t.elem_type);
      numels.push_back(t.numel);
    }
    s.pack(d + "::names", names);
    s.pack(d + "::shapes", shapes);
    s.pack(d + "::elem_types", elem_types);
    s.pack(d + "::numels", numels);
  }
  static void unpack_tensors(DeserializingStream& s, const std::string& d,
                             std::vector<OnnxTensorInfo>& v) {
    std::vector<std::string> names;
    std::vector<std::vector<casadi_int>> shapes;
    std::vector<casadi_int> elem_types, numels;
    s.unpack(d + "::names", names);
    s.unpack(d + "::shapes", shapes);
    s.unpack(d + "::elem_types", elem_types);
    s.unpack(d + "::numels", numels);
    v.clear();
    for (size_t k = 0; k < names.size(); ++k)
      v.push_back(OnnxTensorInfo{names[k], shapes[k], elem_types[k], numels[k]});
  }
 
  void OnnxFunction::serialize_type(SerializingStream &s) const {
    FunctionInternal::serialize_type(s);
    PluginInterface<OnnxFunction>::serialize_type(s);
  }
 
  void OnnxFunction::serialize_body(SerializingStream &s) const {
    FunctionInternal::serialize_body(s);
    s.version("OnnxFunction", 2);
    s.pack("OnnxFunction::model_data", std::string(model_data_.begin(), model_data_.end()));
    pack_tensors(s, "OnnxFunction::in", in_);
    pack_tensors(s, "OnnxFunction::out", out_);
    pack_tensors(s, "OnnxFunction::all_in", all_in_);
    s.pack("OnnxFunction::in_src", in_src_);
    s.pack("OnnxFunction::in_val", in_val_);
    s.pack("OnnxFunction::model_inputs",
           std::vector<std::string>(model_inputs_.begin(), model_inputs_.end()));
    s.pack("OnnxFunction::model_outputs",
           std::vector<std::string>(model_outputs_.begin(), model_outputs_.end()));
    s.pack("OnnxFunction::fwd_dim", fwd_dim_);
    s.pack("OnnxFunction::adj_dim", adj_dim_);
    s.pack("OnnxFunction::input_values", input_values_);
    s.pack("OnnxFunction::model_path", model_path_);
    s.pack("OnnxFunction::dim_bindings", dim_bindings_);
    s.pack("OnnxFunction::input_shapes", input_shapes_);
    s.pack("OnnxFunction::derivative_opts", derivative_opts_);
    s.pack("OnnxFunction::builder_opts", builder_opts_);
  }
 
  OnnxFunction::OnnxFunction(DeserializingStream& s) : FunctionInternal(s) {
    int version = s.version("OnnxFunction", 1, 2);
    std::string bytes;
    s.unpack("OnnxFunction::model_data", bytes);
    model_data_.assign(bytes.begin(), bytes.end());
    unpack_tensors(s, "OnnxFunction::in", in_);
    unpack_tensors(s, "OnnxFunction::out", out_);
    unpack_tensors(s, "OnnxFunction::all_in", all_in_);
    s.unpack("OnnxFunction::in_src", in_src_);
    s.unpack("OnnxFunction::in_val", in_val_);
    std::vector<std::string> mi, mo;
    s.unpack("OnnxFunction::model_inputs", mi);
    model_inputs_ = std::set<std::string>(mi.begin(), mi.end());
    s.unpack("OnnxFunction::model_outputs", mo);
    model_outputs_ = std::set<std::string>(mo.begin(), mo.end());
    s.unpack("OnnxFunction::fwd_dim", fwd_dim_);
    s.unpack("OnnxFunction::adj_dim", adj_dim_);
    s.unpack("OnnxFunction::input_values", input_values_);
    if (version >= 2) {
      s.unpack("OnnxFunction::model_path", model_path_);
      s.unpack("OnnxFunction::dim_bindings", dim_bindings_);
      s.unpack("OnnxFunction::input_shapes", input_shapes_);
      s.unpack("OnnxFunction::derivative_opts", derivative_opts_);
      s.unpack("OnnxFunction::builder_opts", builder_opts_);
    }
  }
 
  ProtoFunction* OnnxFunction::deserialize(DeserializingStream& s) {
    return PluginInterface<OnnxFunction>::deserialize(s);
  }
 
  void OnnxFunction::init(const Dict& opts) {
    for (auto&& op : opts) {
      if (op.first == "fwd_dim") fwd_dim_ = op.second.to_string();
      else if (op.first == "adj_dim") adj_dim_ = op.second.to_string();
      if (op.first == "provider" || op.first == "fwd_dim" || op.first == "adj_dim")
        derivative_opts_[op.first] = op.second;
    }
    // Baked inputs are fed a fixed value, not exposed as Function inputs
    std::vector<OnnxTensorInfo> exposed;
    for (const OnnxTensorInfo& t : in_) if (!input_values_.count(t.name)) exposed.push_back(t);
    in_ = exposed;
    build_io_map();  // baked-feed map over all model inputs
    Dict inferred_opts = opts;
    if (!opts.count("is_diff_in")) {
      std::vector<bool> inferred;
      bool found = false;
      auto infer = [&](const std::string& kind, const std::set<std::string>& inputs,
                       const std::set<std::string>& outputs) -> bool {
        bool signature = false;
        for (const auto& y : out_) {
          signature |= kind == "fwd" ? outputs.count("fwd_" + y.name)
                                     : inputs.count("adj_" + y.name);
        }
        if (!signature) return false;
        std::vector<bool> mask;
        for (const auto& x : in_) {
          mask.push_back(kind == "fwd" ? inputs.count("fwd_" + x.name)
                                       : outputs.count("adj_" + x.name));
        }
        casadi_assert(!found || mask == inferred,
          "ONNX derivative signatures disagree on differentiable inputs; specify is_diff_in");
        inferred = mask;
        found = true;
        return true;
      };
      for (const std::string kind : {"fwd", "adj"}) {
        if (infer(kind, model_inputs_, model_outputs_)) continue;
        std::string path = derivative_path(kind);
        if (path.empty() || !Filesystem::exists(path)) continue;
        GraphBuilder sibling(path, builder_opts_);
        std::set<std::string> inputs, outputs;
        for (const Node& node : sibling.get()->node_list()) {
          (node.io == "input" ? inputs : outputs).insert(node.name);
        }
        infer(kind, inputs, outputs);
      }
      if (found) inferred_opts["is_diff_in"] = inferred;
    }
    FunctionInternal::init(inferred_opts);
  }
 
  void OnnxFunction::build_io_map() {
    // in_src: exposed-arg index (>=0), -2 baked value, or -1 unwired (fed a default later)
    in_src_.clear();
    in_val_.clear();
    for (const OnnxTensorInfo& t : all_in_) {
      casadi_int src = -1;
      for (casadi_int j = 0; j < static_cast<casadi_int>(in_.size()); ++j)
        if (in_[j].name == t.name) { src = j; break; }
      auto bv = input_values_.find(t.name);
      if (src < 0 && bv != input_values_.end()) {
        casadi_assert(static_cast<casadi_int>(bv->second.size()) == t.numel,
                      "Baked value for '" + t.name + "' has " + str(bv->second.size())
                      + " elements, expected " + str(t.numel));
        src = -2;
        for (double v : bv->second) in_val_.push_back(v);
      } else {
        in_val_.insert(in_val_.end(), t.numel, 0.0);  // placeholder block keeps voff aligned
      }
      in_src_.push_back(src);
    }
  }
 
  Sparsity OnnxFunction::tensor_sparsity(const std::vector<casadi_int>& shape) {
    // CasADi is 2-D: rank-0/1/2 map directly, higher ranks flatten to a column vector
    if (shape.empty()) return Sparsity::dense(1, 1);
    if (shape.size() == 1) return Sparsity::dense(shape[0], 1);
    if (shape.size() == 2) return Sparsity::dense(shape[0], shape[1]);
    casadi_int numel = 1;
    for (casadi_int d : shape) numel *= d;
    return Sparsity::dense(numel, 1);
  }
 
  Function OnnxFunction::wrap_derivative(const std::string& kind, const std::string& name,
      const std::vector<std::string>& inames, const std::vector<std::string>& onames,
      const std::vector<Sparsity>& in_sp, const std::vector<Sparsity>& out_sp,
      const Dict& dim_bind, const Dict& opts) const {
    bool embedded = has_derivative(kind, model_inputs_, model_outputs_);
    GraphBuilder b = embedded ? GraphBuilder(name, model_data_, "onnx", builder_opts_)
                             : GraphBuilder(derivative_path(kind), builder_opts_);
    GraphBuilderInternal* bi = b.get();
    bi->model_path_ = derivative_path(kind);
    bi->dim_bindings_ = dim_bindings_;
    for (auto&& d : dim_bind) bi->dim_bindings_[d.first] = d.second;
    bi->input_shapes_ = input_shapes_;
    bi->input_values_ = input_values_;
    std::set<std::string> inputs, outputs;
    for (const Node& n : bi->node_list()) {
      (n.io == "input" ? inputs : outputs).insert(n.name);
    }
    casadi_assert(has_derivative(kind, inputs, outputs),
      "ONNX derivative '" + bi->name_ + "' has an incomplete " + kind + " signature");
    std::vector<std::string> cin, con;
    for (const std::string& nm : inames) if (inputs.count(nm)) cin.push_back(nm);
    for (const std::string& nm : onames) if (outputs.count(nm)) con.push_back(nm);
    // Infer unbound matrix axes, including packed columns such as nadj * nfwd.
    for (const Node& n : bi->node_list()) {
      if (n.dimension.size() != 2) continue;
      const auto& names = n.io == "input" ? inames : onames;
      const auto& sparsities = n.io == "input" ? in_sp : out_sp;
      for (size_t i = 0; i < names.size(); ++i) {
        if (names[i] != n.name) continue;
        for (size_t k = 0; k < 2; ++k) {
          const std::string& param = n.dim_params[k];
          if (n.dimension[k] < 0 && !param.empty() && !bi->dim_bindings_.count(param))
            bi->dim_bindings_[param] = k == 0 ? sparsities[i].size1() : sparsities[i].size2();
        }
      }
    }
    // Preserve differentiability through both the ONNX core and its signature wrapper.
    std::vector<bool> din(inames.size(), true), dout(onames.size(), true);
    auto di = opts.find("is_diff_in"), do_ = opts.find("is_diff_out");
    if (di != opts.end()) {
      din = di->second.to_bool_vector();
    } else {
      for (size_t i = 0; i < in_.size(); ++i) din[i] = diff_in(i);
      for (size_t j = 0; j < out_.size(); ++j) din[in_.size() + j] = diff_out(j);
    }
    if (do_ != opts.end()) {
      dout = do_->second.to_bool_vector();
    } else if (kind == "jac") {
      for (size_t j = 0; j < out_.size(); ++j)
        for (size_t i = 0; i < in_.size(); ++i)
          dout[j * in_.size() + i] = diff_out(j) && diff_in(i);
    }
    std::vector<bool> cdi, cdo;
    for (size_t i = 0; i < inames.size(); ++i) if (inputs.count(inames[i])) cdi.push_back(din[i]);
    for (size_t i = 0; i < onames.size(); ++i) if (outputs.count(onames[i])) cdo.push_back(dout[i]);
    Dict copts = derivative_opts_;
    copts["is_diff_in"] = cdi;
    copts["is_diff_out"] = cdo;
    Function g = OnnxFunction::create(plugin_name(), name + "_core", bi, cin, con, copts);
    for (size_t i = 0; i < inames.size(); ++i) {
      if (!inputs.count(inames[i])) continue;
      casadi_assert(g.sparsity_in(inames[i]).size() == in_sp[i].size(),
        "ONNX derivative input '" + inames[i] + "' has an incompatible shape");
    }
    for (size_t i = 0; i < onames.size(); ++i) {
      if (!outputs.count(onames[i])) continue;
      casadi_assert(g.sparsity_out(onames[i]).size() == out_sp[i].size(),
        "ONNX derivative output '" + onames[i] + "' has an incompatible shape");
    }
    // Present CasADi's full derivative signature: feed present inputs, zero-fill absent outputs
    std::map<std::string, MX> m;
    std::vector<MX> args(inames.size());
    for (size_t i = 0; i < inames.size(); ++i) {
      // Omitted inputs must not keep unused primal evaluations alive.
      Sparsity sp = inputs.count(inames[i]) ? in_sp[i] : Sparsity(in_sp[i].size());
      args[i] = MX::sym(inames[i], sp);
      m[inames[i]] = args[i];
    }
    std::vector<MX> gin;
    for (const std::string& nm : cin) gin.push_back(m[nm]);
    std::vector<MX> gout = g(gin);
    std::map<std::string, MX> mo;
    for (size_t k = 0; k < con.size(); ++k) mo[con[k]] = gout[k];
    std::vector<MX> outs(onames.size());
    for (size_t j = 0; j < onames.size(); ++j) {
      auto it = mo.find(onames[j]);
      outs[j] = it != mo.end() ? it->second : MX::zeros(out_sp[j]);
    }
    Dict wopts;
    wopts["is_diff_in"] = din;
    wopts["is_diff_out"] = dout;
    auto it = opts.find("derivative_of");
    if (it != opts.end()) wopts["derivative_of"] = it->second;
    return Function(name, args, outs, inames, onames, wopts);
  }
 
  std::string OnnxFunction::derivative_path(const std::string& kind) const {
    if (model_path_.empty() || !Filesystem::is_enabled()) return std::string();
    return Filesystem::ensure_trailing_slash(Filesystem::parent_path(model_path_))
      + kind + "_" + Filesystem::filename(model_path_);
  }
 
  bool OnnxFunction::has_derivative(const std::string& kind,
      const std::set<std::string>& inputs, const std::set<std::string>& outputs) const {
    std::string pref = diff_prefix(kind);
    bool any_in = false, any_out = false;
    for (size_t i = 0; i < in_.size(); ++i) {
      if (!diff_in(i)) continue;
      any_in = true;
      if (kind == "fwd" && !inputs.count(pref + in_[i].name)) return false;
      if (kind == "adj" && !outputs.count(pref + in_[i].name)) return false;
      if (kind == "jac") {
        for (size_t j = 0; j < out_.size(); ++j) {
          if (diff_out(j) && !outputs.count("jac_" + out_[j].name + "_" + in_[i].name))
            return false;
        }
      }
    }
    for (size_t j = 0; j < out_.size(); ++j) {
      if (!diff_out(j)) continue;
      any_out = true;
      if (kind == "fwd" && !outputs.count(pref + out_[j].name)) return false;
      if (kind == "adj" && !inputs.count(pref + out_[j].name)) return false;
    }
    return any_in && any_out;
  }
 
  bool OnnxFunction::has_forward(casadi_int nfwd) const {
    return has_derivative("fwd", model_inputs_, model_outputs_)
      || (!model_path_.empty() && Filesystem::is_enabled()
        && Filesystem::exists(derivative_path("fwd")));
  }
 
  Function OnnxFunction::get_forward(casadi_int nfwd, const std::string& name,
                                     const std::vector<std::string>& inames,
                                     const std::vector<std::string>& onames,
                                     const Dict& opts) const {
    std::vector<Sparsity> isp, osp;
    for (const OnnxTensorInfo& x : in_) isp.push_back(tensor_sparsity(x.shape));
    for (const OnnxTensorInfo& y : out_) isp.push_back(tensor_sparsity(y.shape));
    for (const OnnxTensorInfo& x : in_) {
      Sparsity s = tensor_sparsity(x.shape);
      isp.push_back(Sparsity::dense(s.size1(), nfwd * s.size2()));
    }
    for (const OnnxTensorInfo& y : out_) {
      Sparsity s = tensor_sparsity(y.shape);
      osp.push_back(Sparsity::dense(s.size1(), nfwd * s.size2()));
    }
    std::string dim = diff_prefix("fwd");
    dim.replace(0, 3, fwd_dim_);
    dim.pop_back();  // Drop the trailing underscore: nfwd, nfwd2, nfwd3, ...
    Dict db; db[dim] = nfwd;
    return wrap_derivative("fwd", name, inames, onames, isp, osp, db, opts);
  }
 
  bool OnnxFunction::has_reverse(casadi_int nadj) const {
    return has_derivative("adj", model_inputs_, model_outputs_)
      || (!model_path_.empty() && Filesystem::is_enabled()
        && Filesystem::exists(derivative_path("adj")));
  }
 
  Function OnnxFunction::get_reverse(casadi_int nadj, const std::string& name,
                                     const std::vector<std::string>& inames,
                                     const std::vector<std::string>& onames,
                                     const Dict& opts) const {
    std::vector<Sparsity> isp, osp;
    for (const OnnxTensorInfo& x : in_) isp.push_back(tensor_sparsity(x.shape));
    for (const OnnxTensorInfo& y : out_) isp.push_back(tensor_sparsity(y.shape));
    for (const OnnxTensorInfo& y : out_) {
      Sparsity s = tensor_sparsity(y.shape);
      isp.push_back(Sparsity::dense(s.size1(), nadj * s.size2()));
    }
    for (const OnnxTensorInfo& x : in_) {
      Sparsity s = tensor_sparsity(x.shape);
      osp.push_back(Sparsity::dense(s.size1(), nadj * s.size2()));
    }
    std::string dim = diff_prefix("adj");
    dim.replace(0, 3, adj_dim_);
    dim.pop_back();  // Drop the trailing underscore: nadj, nadj2, nadj3, ...
    Dict db; db[dim] = nadj;
    return wrap_derivative("adj", name, inames, onames, isp, osp, db, opts);
  }
 
  bool OnnxFunction::has_jacobian() const {
    return has_derivative("jac", model_inputs_, model_outputs_)
      || (!model_path_.empty() && Filesystem::is_enabled()
        && Filesystem::exists(derivative_path("jac")));
  }
 
  Function OnnxFunction::get_jacobian(const std::string& name,
                                      const std::vector<std::string>& inames,
                                      const std::vector<std::string>& onames,
                                      const Dict& opts) const {
    std::vector<Sparsity> isp, osp;
    for (const OnnxTensorInfo& x : in_) isp.push_back(tensor_sparsity(x.shape));
    for (const OnnxTensorInfo& y : out_) isp.push_back(tensor_sparsity(y.shape));
    for (const OnnxTensorInfo& y : out_) {
      Sparsity so = tensor_sparsity(y.shape);
      for (const OnnxTensorInfo& x : in_) {
        Sparsity si = tensor_sparsity(x.shape);
        osp.push_back(Sparsity::dense(so.size1() * so.size2(), si.size1() * si.size2()));
      }
    }
    return wrap_derivative("jac", name, inames, onames, isp, osp, Dict(), opts);
  }
 
  Function OnnxFunction::create(const std::string& solver, const std::string& name,
                                const GraphBuilderInternal* gb,
                                const std::vector<std::string>& inputs,
                                const std::vector<std::string>& outputs,
                                const Dict& opts) {
    return Function::create(getPlugin(solver).creator(name, gb, inputs, outputs, opts), opts);
  }
 
  Function OnnxFunction::from_model_data(const std::string& solver, const std::string& name,
                                         const std::vector<uint8_t>& model_data,
                                         const std::vector<std::string>& inputs,
                                         const std::vector<std::string>& outputs,
                                         const Dict& opts) {
    // Stage the model + configuration in a GraphBuilder, then freeze an OnnxFunction snapshot
    GraphBuilder b(name, model_data, "onnx");
    GraphBuilderInternal* bi = b.get();
    Dict fopts;
    for (auto&& op : opts) {
      if (op.first == "dim_bindings") {
        for (auto&& d : static_cast<Dict>(op.second)) bi->dim_bindings_[d.first] = d.second;
      } else if (op.first == "input_shapes") {
        for (auto&& d : static_cast<Dict>(op.second))
          bi->input_shapes_[d.first] = d.second.to_int_vector();
      } else if (op.first == "input_values") {
        for (auto&& d : static_cast<Dict>(op.second))
          bi->input_values_[d.first] = d.second.to_double_vector();
      } else {
        fopts[op.first] = op.second;
      }
    }
    return OnnxFunction::create(solver, name, bi, inputs, outputs, fopts);
  }
 
} // namespace casadi