stan-dev / math

The Stan Math Library is a C++ template library for automatic differentiation of any order using forward, reverse, and mixed modes. It includes a range of built-in functions for probabilistic modeling, linear algebra, and equation solving.
https://mc-stan.org
BSD 3-Clause "New" or "Revised" License
753 stars 188 forks source link

All files in prim/fun and prim/constraint should have associated files in rev/fun, rev/constraint and fwd/fun, fwd/constrain #3123

Open SteveBronder opened 1 week ago

SteveBronder commented 1 week ago

Description

When we do not have a reverse or forward mode specialization for a function we default to the function available in prim. When we write functions with no reverse or forward mode specializations we assume the sub functions (or sub sub functions) that make up that larger function will then use reverse or forward mode specializations. Take the log1p_exp example below, Stan has reverse and forward mode for this function but I'm just using it as an example. This log1p function is generically templated so that it would accept double, var, and fvar<T> types.

template <typename T>
inline auto log1p_exp(T a) {
  // like log_sum_exp below with b=0.0; prevents underflow
  if (a > 0.0) {
    return a + log1p(exp(-a));
  }
  return log1p(exp(a));
}

If the forward and reverse mode functions for log1p and exp are not included beforehand, Stan would either

  1. Default to the prim version of the function
  2. (worse case) call the standard library function
  3. (best case) Throw a compiler error

To make sure those includes come in correctly, for each file in prim that does not already have a specialization in rev or fwd we should have a file in the equivalent path of rev/fwd. That file will contain includes for any sub functions used in the prim function that have specializations in rev/(fwd). The final include will include the prim function that will then have the specializations for its sub-functions available to it. This will also simplify our include scheme for fwd/prim/rev as well. The new rules will be that all includes in rev(fwd) must come from rev(fwd), besides a prim implementation of the same file name.

For feature requests:

Given log1p above, it's rev file would look like the following

#include <stan/math/rev/core.hpp>
#include <stan/math/rev/fun/log1p.hpp>
#include <stan/math/rev/fun/exp.hpp>
#include <stan/math/prim/fun/log1p_exp.hpp>

Current Version:

v4.9.0

syclik commented 1 week ago

@SteveBronder, thanks for putting up this issue. I think the issue itself could use a few more details on the technical side in order to be fully specified and implementable.

Before that, a couple of questions:

  1. Is the action of populating all the files what we want to do?
    • It's a lot of work. It's finite and scoped, so it's doable. I'm just putting that out there. (It's fine; a bit of work in exchange for clarity is worth it.)
  2. Does this properly solve the problem of having the incorrect specialization in scope?
    • I think if we're having trouble with calling the wrong instantiation, then this action doesn't actually solve that problem. (I think the problem would be something in stan/math/rev/core.hpp or unexpected has instantiated something improperly.)

To go with your example, if log1p_exp compiles with a Stan type, then all the sub functions (+, *, log1p, exp) have been instantiated with the Stan type. If exp was instantiated with prim version and not the specialized version, that would mean that the specialized version wasn't included by the time log1p_exp was instantiated.

I can see that happening if someone were using Math library by trying to selectively include functions, but if we're including any of the top-level includes such as stan/math/rev.hpp then we shouldn't have this problem by construction. For any rev specialization, this file would be included in rev.hpp and all the functions that have specializations would also be included in the proper order.

I think if we have the problem now, it'll still be a problem if we implemented this change. I think it would come from an improper include somewhere in something like core.hpp, which is absolutely possible.

Have you seen this problem recently?