Modifying external C++ function example to return analytical gradient

Developers
1

Hello,
I’m trying to write a custom C++ function with an analytical gradient that I’d like to call from Stan. I’ve found two relevant examples in the documentation:

First, the “Translating Stan to C++” chapter of the cmdstanr guide provides the following example of a custom C++ function:

namespace bernoulli_model_namespace {
          template <typename T0__>  inline  typename
          boost::math::tools::promote_args<T0__>::type  make_odds(const T0__&
theta, std::ostream* pstream__) {
       return theta / (1 - theta);  }
       }

with a corresponding definition in the Stan file:

functions {
       real make_odds(real theta);
}

How might I modify the C++ function to return an analytical gradient?

Second, say I follow the model in Contributing New Functions to Stan and define a new function in the stan::math namespace:

#include "my_foo.hpp"
#include <stan/math/rev/core.hpp>

namespace stan {
  namespace math {
    vector<double> foo(const vector<double>& x) {
      return bar::foo(x);
    }

    vector<var> foo(const vector<var>& x) {
      vector<double> a = value_of(x);
      double fa = bar::foo(a);
      vector<double> grad_fa = bar::grad_foo(a);  
      return precomputed_gradients(fa, x, grad_fa);
    }
  }
}

How should I call this function foo from a Stan model?

Thanks!

2 Likes
2

Hi, sorry for not getting to you earlier.

I think that you’ll need at least passing familiarity with how C++ templates work. I am not an expert on this, but there are at least two ways you can go about using analytic gradients with custom functions.

As you probably have noticed, the precomputed_gradients returns a (vector of) var (autodiff variable) so you only can use it when the arguments are actually autodiff variables. If your custom function is only ever called with autodiff variables then something like:

functions {
  vector foo(vector x);
}

parameters {
  vector[5] p;
}

model {
  foo(p);
}

with

vector<var> foo(const vector<var>& x, std::ostream* pstream__) {
      vector<double> a = value_of(x);
      double fa = bar::foo(a);
      vector<double> grad_fa = bar::grad_foo(a);  
      return precomputed_gradients(fa, x, grad_fa);
    }

should work (I didn’t test the code). If you need the function to work with both data and parameters, you will need to write multiple specializations, but AFAIK, the example you posted should mostly work, you just need to add the std::ostream* argument.

Alternatively, you can write a more general code that works with both variables and data as described at “Adding new distributions” in the Math docs (Stan Math Library: Probability Distributions)

Hope that gets you started and feel free to ask for additional clarifications.

1 Like