Exposing new distribution in stanc3, with qualified arguments

Developers
1

This is (in some ways) a followup on Subset error with new cmdstanr version.

After concertation with @avehtari, we’ve decided to update the name of the function to include the suffix _lpmf. Nothing too hard: I update the stan-math and the stanc3 code, notably doing an edit to Utils.ml:

let is_distribution_name s =
  (not
     ( String.is_suffix s ~suffix:"_cdf_log"
     || String.is_suffix s ~suffix:"_ccdf_log"
     || String.equal s "laplace_marginal_poisson_log_lpmf"))
  && List.exists
       ~f:(fun suffix -> String.is_suffix s ~suffix)
       distribution_suffices

With these changes, I can compile and run a model that uses the following line

  target += laplace_marginal_poisson_log_lpmf(y | n_samples, ye, K, phi,
                                         x, delta, delta_int, theta_0);

and in the short term, that’s all I need.

What about

y ~ laplace_marginal_poisson_log(n_samples, ye, ...);

I get the following message

Ill-typed arguments to '~' statement. No distribution 'laplace_marginal_poisson_log' was found with the correct signature.

which is a transpiler issue. Fair enough, I didn’t do much to indicate laplace_marginal_poisson_log was a distribution. The relevant edit in Stan_math_signatures.ml should be

let distributions = 
[ (full_lpmf, "beta_binomial", [DVInt; DVInt; DVReal; DVReal])
  ; (full_lpdf, "beta", [DVReal; DVReal; DVReal])

  . . .
 
 ;  ([Lpmf], "lapalce_marginal_poisson_log",
    [ (DataOnly, UArray UInt); (DataOnly, UArray UInt)
  ; ( AutoDiffable
    , UFun
        ( [ (AutoDiffable, UVector); (DataOnly, UArray UVector)
          ; (DataOnly, UArray UReal); (DataOnly, UArray UInt) ]
        , ReturnType UMatrix ) )
  ; (AutoDiffable, UVector); (DataOnly, UArray UVector)
  ; (DataOnly, UArray UReal); (DataOnly, UArray UInt)
  ; (AutoDiffable, UVector) ]) ]

which prompts the error message

MacBook-Pro:stanc3 charlesm$ dune runtest test/integration
File "src/middle/Stan_math_signatures.ml", line 192, characters 5-9:
Error: This variant expression is expected to have type fkind list
       The constructor Lpmf does not belong to type list
MacBook-Pro:stanc3 charlesm$ dune runtest test/integration
File "src/middle/Stan_math_signatures.ml", line 193, characters 6-29:
Error: This expression has type 'a * 'b
       but an expression was expected of type dimensionality

I’m not sure what the error is here. I also assume there are some additional moving parts from using qualifier arguments. Any help would be very much welcomed.

@rok_cesnovar

2

EDIT: ignore this

I need to look into the problem a bit more, but the first thing that comes to mind is that

y ~ laplace_marginal_poisson_log(n_samples, ye, ...);

is actually transformed to

target += laplace_marginal_poisson_log_propto_lpmf(y | n_samples, ye, ...);

The propto infix is then used at the code-generating end to differentiate between normalized and unnormalized lpdfs/lpmfs.

So you also need to do:

let is_distribution_name s =
  (not
     ( String.is_suffix s ~suffix:"_cdf_log"
     || String.is_suffix s ~suffix:"_ccdf_log"
     || String.equal s "laplace_marginal_poisson_log_lpmf"
     || String.equal s "laplace_marginal_poisson_log_propto_lpmf"))
  && List.exists
       ~f:(fun suffix -> String.is_suffix s ~suffix)
       distribution_suffices

If you are however already changing the math side, I would advise just adding the propto template argument for the sake of consistency with other distributions. You can just ignore it if there is nothing to be gained with propto=true. Which would then mean you do not have to make an exception for this new function.

3

Also a typo here, should be laplace_

1 Like
4

Actually @charlesm93 ignore my first comment. It is related to _propto_, but that solution wont work.

The logic to handle when to generate template lives here
I added a few comments to explain what is going on.

let demangle_propto_name udf f =
  (* exceptions with _log; handle them as regular functions*)
  if f = "multiply_log" || f = "binomial_coefficient_log" then f
  (* is the distribution proportional (has _propto_)? remove _propto_ and add the template parameter (propto__ = true) *)
  else if Utils.is_propto_distribution f then
    Utils.stdlib_distribution_name f ^ "<propto__>"
  else if
  (* is this a normalized distribution (not proportional) ? add the template parameter, hardcoded to false  *)
    Utils.is_distribution_name f || (udf && (is_user_dist f || is_user_lp f))
  then f ^ "<false>"
  (* None of the above? just use the name *)
  else f

My suggestion in order to make the prototype work is to add another exception

 if f = "multiply_log" || f = "binomial_coefficient_log" then  f
  else if Stan_math_signatures.is_laplace_fn f then
    Utils.stdlib_distribution_name f
  else if Utils.is_propto_distribution f then
    Utils.stdlib_distribution_name f ^ "<propto__>"
  else if
    Utils.is_distribution_name f || (udf && (is_user_dist f || is_user_lp f))
  then f ^ "<false>"
  else f

I put my suggestions on a branch, see the comparison here: https://github.com/stan-dev/stanc3/compare/try-laplace_approximation2...laplace-approx-suggestions-rok

With these changes the test models compile to c++ as well as executables with the math branch set.

I would however, suggest to add the propto template parameter in Math before doing the Math/stanc3 PRs. All lpdfs/lpmfs have it and you would avoid doing this exceptions. We would really like to avoid these kind of exceptions in stanc3, especially if the fix is simple.