Method export_with_callback

mlx/export.cpp:698–756

};

void FunctionExporter::export_with_callback(
    const std::vector<array>& inputs,
    const std::vector<array>& outputs,
    const std::vector<array>& tape,
    const std::vector<std::string>& kwarg_keys) {
  NodeNamer namer{};
  auto to_vector_data = [&namer](const auto& arrays) {
    std::vector<std::tuple<std::string, Shape, Dtype>> data;
    for (auto& a : arrays) {
      data.emplace_back(namer.get_name(a), a.shape(), a.dtype());
    }
    return data;
  };

  // Callback on the inputs
  callback({{"type", "inputs"}, {"inputs", to_vector_data(inputs)}});
  std::vector<std::pair<std::string, std::string>> keyword_inputs;
  for (int i = inputs.size() - kwarg_keys.size(), j = 0; i < inputs.size();
       ++i, ++j) {
    keyword_inputs.emplace_back(kwarg_keys[j], namer.get_name(inputs[i]));
  }
  callback({{"type", "keyword_inputs"}, {"keywords", keyword_inputs}});

  // Callback on the outputs
  callback({{"type", "outputs"}, {"outputs", to_vector_data(outputs)}});

  // Callback on the constants
  {
    std::unordered_set<std::uintptr_t> input_set;
    for (auto& in : inputs) {
      input_set.insert(in.id());
    }
    std::vector<std::pair<std::string, array>> new_constants;
    for (auto& arr : tape) {
      if (arr.has_primitive() || input_set.find(arr.id()) != input_set.end()) {
        continue;
      }
      if (constants.insert({arr.id(), arr}).second) {
        new_constants.emplace_back(namer.get_name(arr), arr);
      }
    }
    callback({{"type", "constants"}, {"constants", new_constants}});
  }
  auto factory = PrimitiveFactory();

  // Callback for each primitive in the tape
  for (auto& arr : tape) {
    if (!arr.has_primitive()) {
      continue;
    }
    auto [name, state] = factory.extract_state(arr.primitive_ptr());
    callback(
        {{"type", "primitive"},
         {"inputs", to_vector_data(arr.inputs())},
         {"outputs", to_vector_data(arr.outputs())},
         {"name", name},
         {"arguments", state}});
  }