Function py_value_and_grad

python/src/transforms.cpp:69–307

}

auto py_value_and_grad(
    const nb::callable& fun,
    std::vector<int> argnums,
    std::unordered_set<std::string> argnames,
    const std::string& error_msg_tag,
    bool scalar_func_only) {
  // Sanitize argnums
  if (argnums.size() == 0 && argnames.size() == 0) {
    throw std::invalid_argument(
        error_msg_tag + " Gradient wrt no argument requested");
  }
  for (auto arg : argnums) {
    std::sort(argnums.begin(), argnums.end());
    if (argnums[0] < 0) {
      std::ostringstream msg;
      msg << error_msg_tag
          << " Can't compute the gradient of negative argument index "
          << argnums[0];
      throw std::invalid_argument(msg.str());
    }
    for (int i = 1; i < argnums.size(); ++i) {
      if (argnums[i] == argnums[i - 1]) {
        std::ostringstream msg;
        msg << error_msg_tag << " Duplicate argument index " << argnums[0]
            << " is not allowed.";
        throw std::invalid_argument(msg.str());
      }
    }
  }

  return [fun, argnums, argnames, error_msg_tag, scalar_func_only](
             nb::args& args, nb::kwargs& kwargs) {
    // Sanitize the input
    if (argnums.size() > 0 && argnums.back() >= args.size()) {
      std::ostringstream msg;
      msg << error_msg_tag << " Can't compute the gradient of argument index "
          << argnums.back() << " because the function is called with only "
          << args.size() << " positional arguments.";
      throw std::invalid_argument(msg.str());
    }

    for (auto& key : argnames) {
      if (!kwargs.contains(key)) {
        std::ostringstream msg;
        msg << error_msg_tag
            << " Can't compute the gradient of keyword argument '" << key
            << "' because the function is called with the "
            << "following keyword arguments {";
        for (auto item : kwargs) {
          msg << nb::cast<std::string>(item.first) << ",";
        }
        msg << "}";
        throw std::invalid_argument(msg.str());
      }
    }

    // Collect the arrays
    std::vector<mx::array> arrays;