Method export_function

mlx/export.cpp:758–877

}

void FunctionExporter::export_function(const Args& args, const Kwargs& kwargs) {
  if (closed) {
    throw std::runtime_error(
        "[export_function] Attempting to write after exporting is closed.");
  }
  auto sorted_kwargs =
      std::map<std::string, array>(kwargs.begin(), kwargs.end());
  auto [fentry, inserted] = ftable->emplace(args, sorted_kwargs);
  if (!inserted) {
    throw std::runtime_error(
        "[export_function] Attempting to export a function twice with "
        "the same signature is not allowed.");
  }

  // Flatten the inputs to the function for tracing
  std::vector<std::string> kwarg_keys;
  auto inputs = args;
  for (auto& [k, v] : sorted_kwargs) {
    kwarg_keys.push_back(k);
    inputs.push_back(v);
  }

  auto flat_fun = [this, &kwarg_keys](const Args& flat_args) {
    auto args = Args(flat_args.begin(), flat_args.end() - kwarg_keys.size());
    Kwargs kwargs;
    auto it = flat_args.end() - kwarg_keys.size();
    ;
    for (auto& k : kwarg_keys) {
      kwargs.insert({k, *it++});
    }
    return detail::ArraysAndExtra{fun(args, kwargs), nullptr};
  };

  // Trace to build the graph
  auto [trace_inputs, trace_outputs, extra] =
      detail::compile_trace(flat_fun, inputs, ftable->shapeless);

  // DFS the graph and get the tape
  auto [tape, parents_map] =
      detail::compile_dfs(trace_inputs, trace_outputs, inputs);

  detail::compile_simplify(tape, parents_map, trace_outputs, /* passes */ 3);

  // Update the table entry
  fentry.kwarg_keys = kwarg_keys;
  fentry.inputs = trace_inputs;

  count++;

  if (callback) {
    export_with_callback(trace_inputs, trace_outputs, tape, kwarg_keys);
    return;
  }

  // Update the header
  auto pos = os.tell();
  os.seek(0);
  write_header(os, count, ftable->shapeless);