Skip to content

Conversation

@SteveBronder
Copy link
Collaborator

@SteveBronder SteveBronder commented Nov 1, 2022

Summary

This creates a class propagate_partials that is a drop in replacement for operands_and_partials. A new class was created so that stan math can stay backwards compatible. Underneath the hood propagate_partials is very similar to operands_and_partials, but now each edge is stored inside of a tuple. This has a small benefit of saving some memory relative to operands_and_partials since we previously allocated edges for 5 operations even if we only needed one. The main benefit is that this can support any number of operations and is a drop in replacement for operands_and_partials.

This change also makes for nicer code, where previously we had to include a lengthy template for operands_and_partials we can now just use auto and call the make_partials_propagator() function which then creates and returns the appropriate partials_propagator class. This is similar to the standard libraries make_tuple() function.

  // before
  operands_and_partials<decltype(y_col), decltype(mu_col), decltype(kappa_col)>
      ops_partials(y_col, mu_col, kappa_col);
  // now
  auto ops_partials = make_partials_propagator(y_col, mu_col, kappa_col);

Instead of operands_and_partials's edge1_ and edge2_ the partials_propagator class has an associated function called edge() which is used in the same way as std::tuple's std::get() function

// before
ops_partials.edge1_.partials_ = elt_divide(-P, Pi_cl);
// now
edge<0>(ops_partials).partials_ = elt_divide(-P, Pi_cl);

Tests

Since this is a drop in replacement for operands_and_partials I took all of the tests for that classe and wrote the corresponding ones for partials_propagator. They can be run via

./runTests.py -j2 ./test/unit/math/ -f partials_propagator

Side Effects

I changed all of the distributions to use make_partial_propagator so a lot of files are changed. I did this because I wanted to make sure everything worked and the easiest way was to swap them out for one another.

imo the files that really need looked at thoroughly are the operands_and_partials and partial_propagator ones in prim/fwd/rev. I'll double check the distributions match correctly.

Release notes

Allow for for N-ary partials for distributions

Checklist

@SteveBronder
Copy link
Collaborator Author

Staring at the code for using edge() and wondering if there should just be a function for directly getting the partials? So instead of

edge<0>(ops_partials).partials_ = elt_divide(-P, Pi_cl);

It's possible to just call

partials<0>(ops_partials) = elt_divide(-P, Pi_cl);

@SteveBronder
Copy link
Collaborator Author

@andrjohns would you mind taking a look at this? I've looked over all of the _lpdf/_lpmf files so I think if you look over the actual partial_propagator files and operands_and_partials files that would be good.

Not super high priority. This is mostly a nice cleanup. There may be some small optimizations for cases with only 1 or 2 parameters since we only store as many edges as we need now (and then only do reverse passes for those 1 or 2 parameters as well). Though I think we would only see a noticeable speedup when applying the scalar form of the lpdf many times.

Copy link
Collaborator

@andrjohns andrjohns left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You had me scared with the size of the diff, but overall the changes are pretty clean and minor! Just a few comments/suggestions

Copy link
Collaborator

@andrjohns andrjohns left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All LGTM! Feel free to merge if you don't have any other changes

@SteveBronder SteveBronder merged commit e45480c into develop May 11, 2023
@WardBrian WardBrian deleted the feature/tuple-ops-partials3 branch August 5, 2024 00:15
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.

5 participants