Filtering write_array, log_prob, constrain_params, unconstrain_params by parameter in the model class

Developers
1

There’s been some offline conversation about how to best approach the desire to choose which parameters to focus on for these three operations on a Stan model class after it’s been constructed with the data:

  1. Unconstrain parameters individually (but sufficient to do all at once and let the language interface pick out which is which). Currently handled by a method called transform_inits that does all params at once in and out.
  2. select which parameters are written out by write_array (which also serves as a constrain_parameters API call, essentially).
  3. select which parameters to differentiate w.r.t…

I have been proposing these three related type signatures for the API:

template <typename T>
using param_map = std::unordered_map<std::string, std::vector<T>>; // param name -> flattened values

const param_map<double> unconstrain_parameters(const param_map<double>& constrained_params);

const param_map<double> write_array(RNG rng, const param_map<double>& unconstrained_params, 
                                // which param | tparam | gq vars you want written out
                                const std::set<string>& desired_output_vars, 
                                std::ostream* pstream = 0);

template <typename T>
T log_prob(const param_map<T>& autodiff_params, 
           const param_map<double>& fixed_params,
           std::ostream* pstream = 0);

What do folks think about this for solving the problem? There have been some concerns voiced already that might need further research or benchmarking:

  • Can we write these methods in a way that is fast enough? Particular concerns around log_prob and write_array as those are called most often. We could offer both the existing methods and these new overrides or benchmark and see if it’s an issue.
  • We will have to issue runtime errors if insufficient parameters are passed in to calculate the desired quantities and parameters.
  • “We’d need some fancy analysis and code gen.” I think we live in that world now with stanc3 :) I think it would be sufficient to topologically sort the parameters and their dependencies and find any linearization of the graph, variables with 0 dependencies first. Then we just have an if block for each variable that checks if it is needed (desired or needed for another desired variable) and if so, checks that its dependencies have been met, and if so reads it in. If it’s needed but can’t be read, we throw an exception.
  • We could also just attempt to generate them all and then selectively emit, which would still be a win from a memory-us and language-interface perspective.

There was another proposal from @Bob_Carpenter back in the day but it needs to be updated to deal with tuples and other recursive types.

@bgoodri suggested using a boolean array to indicate which vars to write out with write_array. It seems like this could be more efficient and granular at the expense of offloading the work to the interface, and a boolean array doesn’t work for the other 3 methods. Something to consider.

1 Like
2

That doesn’t include any filtering. I was mainly trying to address programatically accessing parameter types, shapes, sizes, and values through C++.

I don’t think tuples are such a big deal for our current string based I/O, in that they can still be indexed. That is, if you have y = (a, b, c), then y.1, y.2 and y.3 pick out a, b, and c and you can keep indexing from there. The problem’s when we want to return things with structure, as that means some encoding like JSON or some kind of variant type, which is a huge pain in C++, or maybe something intermediate.

3

Fair enough.