Rowwise/Colwise Framework

[andrjohns]

This design document outlines a proposed rowwise and colwise frameworks, allowing users to specify custom row- and column-wise functions in Stan. The frameworks allow for a variadic number of iterated and shared arguments.

Rendered markdown is here

Relevant Math branch is here

[t4c1]

Interesting Idea. I have a few questions.

• Can this also support built in functions - ones implemented in Stan Math, or just user defined function?
• What are the pros and cons of the proposed implementation compared to adding support for Eigen::VectorwiseOp to our functions (the return type of .rowwise() and .colwise() functions)?

[rok-cesnovar]

Nice. I am guessing this could then also be parallelized with TBB and thus be a replacement of map_rect with a nicer API?

[andrjohns]

Can this also support built in functions - ones implemented in Stan Math, or just user defined function?

As in the Stan syntax would be: rowwise(input1, log_sum_exp) or similar? I don't know if the transpiler supports using functions in that way

What are the pros and cons of the proposed implementation compared to adding support for Eigen::VectorwiseOp to our functions (the return type of .rowwise() and .colwise() functions)?

That was my idea, but to be perfectly honest I couldn't figure out a way to do it. The VectorwiseOp class seems pretty specialised towards the specific functions, I couldn't see a way to implement it in a general fashion. But if you've got the time to have a look and have some general ideas then I'm happy chase them down.

Nice. I am guessing this could then also be parallelized with TBB and thus be a replacement of map_rect with a nicer API?

Unfortunately not yet, since there's still the same issues with parallel reverse mode which would need to be figured out first

[t4c1]

I don't know if the transpiler supports using functions in that way

I am pretty sure it does not. It was meant more like could we add support for that - I think the Math functions would just need to get wrapped into functors. Althought that question is probably more for @rok-cesnovar.

That was my idea, but to be perfectly honest I couldn't figure out a way to do it.

I guess we would need to add support for VectorwiseOp to all Math and user defined functions, which would be much more work than what you proposed. Although It might be a bit more performant. It is probably not worth it.

[rok-cesnovar]

I am pretty sure it does not. It was meant more like could we add support for that - I think the Math functions would just need to get wrapped into functors.

Indeed. Its not supported right now, but I think generating functors for any math function used in such cowwise/rowwise should not be a big deal if this is something we think would be beneficial.

Unfortunately not yet, since there's still the same issues with parallel reverse mode which would need to be figured out first

Ok, thanks. This is great even without that.

[andrjohns]

Indeed. Its not supported right now, but I think generating functors for any math function used in such cowwise/rowwise should not be a big deal if this is something we think would be beneficial.

Well, something that could get in the way of that, and possibly the framework in general, is that the location of the functor isn't fixed between calls. In other words, the functor could be the second argument:

rowwise(input1, log_sum_exp)

Or the third (or Nth):

rowwise(input1, input2, custom_func, input3)

Is this likely to be problematic with Stanc3? Are the signatures fine to just be:

matrix rowwise(T x, ...)
vector rowwise(T x, ...)

[rok-cesnovar]

No, I dont think that is a problem Stanc3 wise. It will require some additional logic for semantic checking that the types match, but right now I dont see it that being a problem if the C++ side of this is handled as you propose.

In a chat with @spinkney he pointed out that rowwise would also work well with ragged arrays once those are introduced. I think that is also a great + for this.

[andrjohns]

Oh I hadn't thought about ragged arrays, not too familiar with that side of things. I'll take your word for it!

[SteveBronder]

If we had a tuple type in the language (Not sure about the status of that) we could have something like

colwise(functor, tuple(iterated_args), tuple(shared_args));

Then we wouldn't have to worry about that. One thing we could do for a simpler first impl is just support unary functions to iterate over matrices

[rok-cesnovar]

That is a very good point! Tuples make this much easier to do.

@rybern any status updates on tuples in stanc3?

I believe its stan-dev/stanc3#670 first then stan-dev/stanc3#675

[rybern]

@rybern any status updates on tuples in stanc3?

Update is that I finally have some time again! My goal for this week is getting it to a point where I can rope in some other contributors.

[andrjohns]

Oh yeah tuples would majorly simplify this, both on the Stan and C++ side of things. Great news Ryan!!

[andrjohns]

Just revisiting this project now, is it worth waiting for tuples in the language or should I go ahead with the spec as-is? For simplicity, I could just implement the variadic case now, and then add a signature for tuples later when they're available

[SteveBronder]

Let me see if we can get a better eyeball on tuples this week. Right now we are waiting on stan-dev/stan#3036 but after that we should have a better ballpark on time.

One quick thing about the current impl, idt it would work if someone passed a tuple? Though don't think it would require a bit change

template <typename T>
using is_stan_type = disjunction<is_stan_scalar<T>, is_container<T>>;

I think we need something for type detection there that is more than detecting what it's not. Though I'm not sure if there is some type trait you can use to detect a lambda or functor

[andrjohns]

Great! Happy to wait, just wanted to get a better idea on the timeline.

Yeah I went down a bit of a rabbithole but I couldn't find a reliable type trait for detecting both lambdas and functors, so I had to settle on the disjunction approach. Not at all married to it, so any other suggestions are also vv welcome