Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Allow distributions to return vectors #2751

Open
wants to merge 8 commits into
base: develop
Choose a base branch
from

Conversation

SteveBronder
Copy link
Collaborator

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
Copy link
Collaborator Author

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
Copy link
Contributor

betanalpha commented Jun 20, 2022 via email

@bob-carpenter
Copy link
Contributor

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
Copy link
Collaborator Author

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
Copy link
Contributor

betanalpha commented Aug 23, 2022 via email

@bob-carpenter
Copy link
Contributor

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
Copy link
Collaborator Author

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
Copy link
Contributor

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
Copy link
Contributor

betanalpha commented Sep 23, 2022 via email

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

vector-output log density functions
5 participants