[WIP] Allow distributions to return vectors

[SteveBronder]

Summary

This is a prototype for discussion about allowing our distributions to return back vectors instead of just scalars. I only did normal_lpdf as an example of how this would be implemented for the other functions. My goal was to try to reduce the amount of possible code duplication while still keeping our optimizations for the case where we are returning back a scalar type.

The main components here are

1. A new template parameter ReturnType to the lpdf that is either ProbReturnType::Scalar or ProbReturnType::Vector
2. a new class prob_reducer* that for ProbReturnType::Scalar a scalar or for ProbReturnType::Vector holds an Eigen array. For the scalar case, things like += etc. will do summation when the rhs is an Eigen type. For ProbReturnType::Vector when operators such as += see a scalar they propogate the scalar over the entire return Eigen array.

Getting this to work nicely does require a little bit of oddities, for instance in normal_lpdf we previously had

return_t logp = sum(-0.5 * y_scaled_sq);

But now we may want the rhs not summed, but y_scaled_sq could also just be a scalar. So now we sometimes need to construct logp with an explicit constructor telling it the size we want.

prob_return_t<RetType, T_partials_return> logp(-0.5 * y_scaled_sq, N);

So that if y_scaled_sq is a scalar it propogates that for each value in the return matrix of size N

• prob_reducer is probably a bad name and a new one is very welcome. I would like a name that reflects that the class either holds a scalar and then does summations over vector inputs or is a vector which then just returns a vector at the end. Maybe conditional_reducer?

Tests

After stealing a few things from the prob tests I was able to get a version working with the expect_ad test framework. A new function expect_ad_distribution applies the user input over all the possible combinations of types that can go into the probability distributions.

You can see an example of this in the mix test for the normal distribution

./runTests.py test/unit/math/mix/prob/normal_test.cpp

Side Effects

I'm not sure about side effects necessarily but this is a big change that we should be very careful with.

Release notes

Adds framework for probability distributions to return back scalars or Eigen vectors.

Checklist

• Math issue #(issue number) (I couldn't find the issue but I know people have been requesting this forever and there is an issue somewhere)

• Copyright holder: Steve Bronder

  The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
  - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
  - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

• the basic tests are passing

  • unit tests pass (to run, use: ./runTests.py test/unit)
  • header checks pass, (make test-headers)
  • dependencies checks pass, (make test-math-dependencies)
  • docs build, (make doxygen)
  • code passes the built in C++ standards checks (make cpplint)

• the code is written in idiomatic C++ and changes are documented in the doxygen

• the new changes are tested

[SteveBronder]

While I can help get in a framework for this, idt I have time to do this for all of the distributions. Though I do think this could be a cool GSOC project for a summer

[betanalpha]

Is there any proposal for how this feature would be exposed in the language?

…

On Jun 3, 2022, at 1:43 PM, Steve Bronder ***@***.***> wrote: Summary This is a prototype for discussion about allowing our distributions to return back vectors instead of just scalars. I only did normal_lpdf as an example of how this would be implemented for the other functions. My goal was to try to reduce the amount of possible code duplication while still keeping our optimizations for the case where we are returning back a scalar type. The main components here are A new template parameter ReturnType to the lpdf that is either ProbReturnType::Scalar or ProbReturnType::Vector a new class prob_reducer* that for ProbReturnType::Scalar a scalar or for ProbReturnType::Vector holds an Eigen array. For the scalar case, things like += etc. will do summation when the rhs is an Eigen type. For ProbReturnType::Vector when operators such as += see a scalar they propogate the scalar over the entire return Eigen array. Getting this to work nicely does require a little bit of oddities, for instance in normal_lpdf we previously had return_t logp = sum(-0.5 * y_scaled_sq); But now we may want the rhs not summed, but y_scaled_sq could also just be a scalar. So now we sometimes need to construct logp with an explicit constructor telling it the size we want. prob_return_t<RetType, T_partials_return> logp(-0.5 * y_scaled_sq, N); So that if y_scaled_sq is a scalar it propogates that for each value in the return matrix of size N prob_reducer is probably a bad name and a new one is very welcome. I would like a name that reflects that the class either holds a scalar and then does summations over vector inputs or is a vector which then just returns a vector at the end. Maybe conditional_reducer? Tests After stealing a few things from the prob tests I was able to get a version working with the expect_ad test framework. A new function expect_ad_distribution applies the user input over all the possible combinations of types that can go into the probability distributions. You can see an example of this in the mix test for the normal distribution ./runTests.py test/unit/math/mix/prob/normal_test.cpp Side Effects I'm not sure about side effects necessarily but this is a big change that we should be very careful with. Release notes Adds framework for probability distributions to return back scalars or Eigen vectors. Checklist Math issue #(issue number) (I couldn't find the issue but I know people have been requesting this forever and there is an issue somewhere) Copyright holder: Steve Bronder The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses: - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause <https://opensource.org/licenses/BSD-3-Clause>) - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/ <https://creativecommons.org/licenses/by/4.0/>) the basic tests are passing unit tests pass (to run, use: ./runTests.py test/unit) header checks pass, (make test-headers) dependencies checks pass, (make test-math-dependencies) docs build, (make doxygen) code passes the built in C++ standards <https://github.com/stan-dev/stan/wiki/Code-Quality> checks (make cpplint) the code is written in idiomatic C++ and changes are documented in the doxygen the new changes are tested You can view, comment on, or merge this pull request online at: #2751 <#2751> Commit Summary 0890920 <0890920> adds internals needed for supporting vector returning lpdfs d888a0f <d888a0f> get vector working for prim and reverse mode 793b4f8 <793b4f8> prototype reducer to allow distributions to return vectors instead of scalars f7b6c3d <f7b6c3d> clang format File Changes (12 files <https://github.com/stan-dev/math/pull/2751/files>) M stan/math/fwd/functor/operands_and_partials.hpp <https://github.com/stan-dev/math/pull/2751/files#diff-4284c4afdb6b17073e20a9dd03d717a8e3bbf7d3bf5d16b6141fb79fcd5d51f3> (38) M stan/math/prim/functor/operands_and_partials.hpp <https://github.com/stan-dev/math/pull/2751/files#diff-427bca5e937d2310031622d636ea904b7e8de678ca3b6cc5a6c6239b94a66794> (4) A stan/math/prim/functor/prob_reducer.hpp <https://github.com/stan-dev/math/pull/2751/files#diff-086064014ae2e658c7d3beef97afdb2b47f4966b8dca084ca02242d0d19e1456> (250) M stan/math/prim/prob/normal_log.hpp <https://github.com/stan-dev/math/pull/2751/files#diff-b9a2fe6dab2e77a07b11361db8ba0350fce53949376bbf6f841a625517133861> (8) M stan/math/prim/prob/normal_lpdf.hpp <https://github.com/stan-dev/math/pull/2751/files#diff-9f8dd7883399158bc20b615526c4ea7e366d46a2d5a9d8997d6a843e23abf016> (36) M stan/math/rev/functor/operands_and_partials.hpp <https://github.com/stan-dev/math/pull/2751/files#diff-2fdf4c9b8c564ea14ce7ec466abebcd07e33db04b33916d59a5a1dcb970e47a4> (95) A test/unit/math/fwd/prob/normal_vectorized_test.cpp <https://github.com/stan-dev/math/pull/2751/files#diff-d8d80089f031e767e7425c760b1945e131b1c3c5e60653c8a852374090188db2> (11) M test/unit/math/mix/prob/normal_test.cpp <https://github.com/stan-dev/math/pull/2751/files#diff-fa8258f6a77d7078034977ae6fdeea9d98fa3da208b818e70eea5ef7091960a0> (17) A test/unit/math/mix/prob/test_distribution_ad.hpp <https://github.com/stan-dev/math/pull/2751/files#diff-32258f0222d6ba178785d13f1977d3c9603c2648b1dbd24895ef3ff222d2936d> (282) M test/unit/math/prim/prob/normal_log_test.cpp <https://github.com/stan-dev/math/pull/2751/files#diff-f534dfd64ee37336c501403ea21d59f768f488b9b6ef7dd3dd85312c36deeaef> (20) M test/unit/math/rev/prob/normal_log_test.cpp <https://github.com/stan-dev/math/pull/2751/files#diff-d91e8f2f20d99c89105f955133ba691606cca1165185f3c2496d8655efa3ae29> (10) M test/unit/math/test_ad.hpp <https://github.com/stan-dev/math/pull/2751/files#diff-f31332a6f8df42c58e14ab21dab7956a8e29f2d099a14bd1495d4086b88af03a> (1) Patch Links: https://github.com/stan-dev/math/pull/2751.patch <https://github.com/stan-dev/math/pull/2751.patch> https://github.com/stan-dev/math/pull/2751.diff <https://github.com/stan-dev/math/pull/2751.diff> — Reply to this email directly, view it on GitHub <#2751>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AALU3FUJKQOR7AXDBIJSBE3VNI73PANCNFSM5XZWCNSQ>. You are receiving this because you are subscribed to this thread.

[bob-carpenter]

Is there any proposal for how this feature would be exposed in the language?

We'd just use another decorator, like _vec, as in normal_lpdf_vec. I don't think we want to start decorating function calls with special syntax and this isn't an annotation thing as it affects the language.

An alternative would be a general vectorization function:

T0 f(T1, T2, ..., TN)

array[] T0 vectorize(f, (T1 | array[] T1) xs1, array[] (T2 | array[] T2) xs2, ..., (TN | array[] TN) xsN);

Then the use would be

vector[N] lp = vectorize(normal_lupdf, y, mu, sigma);

where now y, mu, and sigma can be scalars or arrays (or maybe vectors, too).

The implementation can just create a view of each argument then apply f in a loop. There's no easy way to do this efficiently, though, as we can't share evaluations like log(sigma) in vectorizing normal_lpdf.

[SteveBronder]

I personally just like having _vec_lpdf as the compiler can pretty easily check if the function ends with _vec_lpdf or _vec_lpmf and apply the approriate templating for returning a vector

[betanalpha]

We’re kind of in an awkward corner with the conventions that density functions collapse the broadcast evaluation by default. To me “vectorization” has always been a bit confusion and something like “nosum” or “nocollapse” or “expanded” would be more explicit, but given the existing conventions “vec_lpdf” or “lpdf_vec” wouldn’t be the worst.

…

On Jul 8, 2022, at 12:20 PM, Steve Bronder ***@***.***> wrote: I personally just like having _vec_lpdf as the compiler can pretty easily check if the function ends with _vec_lpdf or _vec_lpmf and apply the approriate templating for returning a vector — Reply to this email directly, view it on GitHub <#2751 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AALU3FQXB2SUC5RCMH5TOBTVTBILZANCNFSM5XZWCNSQ>. You are receiving this because you commented.

[bob-carpenter]

I agree that "vectorization" is ridiculously overloaded. It technically refers to how the CPU pipelines operations in SSE, AVX, etc. Our general use for functions like exp(vector) is already a stretch in that we're really doing what is called "type lifting" in the programming languages literature---we just map the function elementwise in this case.

For what it's worth, I prefer lpdf_vec, but have no idea why. Either would work.

[SteveBronder]

Yeah it's weird and hard to name. Even vec seems weird. What about for_each? Like normal_lpdf_for_each(...)? I'm also fine with lpdf_vec

[bob-carpenter]

To me, _for_each sounds like it's saying the behavior is a loop. While that's technically true and therefore the name makes sense, I find it too long and not quite evocative enough of the different return type.

Something like _ind or _item would make sense to return item-level densities, but I don't think they're as clear as _vec.

[betanalpha]

To me, _for_each sounds like it's saying the behavior is a loop. While that's technically true and therefore the name makes sense, I find it too long and not quite evocative enough of the different return type. Something like _ind or _item would make sense to return item-level densities, but I don't think they're as clear as _vec.

While “for_each” might not be my favorite if we were redesigning the language I actually kind of like it in the context of the existing language. In particular I think that most users, especially those with less formal programming experience, would interpret “vectorization/broadcasting/lifting to a product space” as a for loop, making “for each” about as intuitive as possible. In any case this is a big enough decision that it’s probably worth a forum discussion to elicit more feedback from users.