Delegate ternary operator to Math `if_else()`

[andrjohns]

At the moment the ternary operator in stan directly transpiles to the C++ ternary operator. However, there have been requests in the past for a vectorised ternary operator:

• https://discourse.mc-stan.org/t/ternary-operator-for-vectors/23149
• https://discourse.mc-stan.org/t/elementwise-ternary-operator/1598

I'm currently working on a ternary vectorisation framework which could be used to vectorise the if_else() function in the Math library: https://github.com/stan-dev/math/issues/2638

If stanc3 were to transpile the ternary operator to if_else, then Stan would automatically support vectorised ternary statements once the if_else function is vectorised.

[bob-carpenter]

Thanks for popping this up, @andrjohns.

The ternary operator in Stan pretty much has to compile to the ternary operator. That's the only way to get the appropriate short-circuiting behavior without some sort of elaborate lazy evaluation wrapper around the arguments. We can extend the ternary operator to vectors, but have to be careful to compile it to a different function. I would recommend actually writing an if_else() function that does not short circuit.

I find the implementation choice in R to be perplexing in this instance.

ifelse(TRUE, c(1, 2), c(3, 4))
[1] 1

The doc is here:

‘ifelse’ returns a value with the same shape as ‘test’ which is filled with elements selected from either ‘yes’ or ‘no’ depending on whether the element of ‘test’ is ‘TRUE’ or ‘FALSE’.

I suppose that's what you get when you conflate scalars and singleton arrays.

I would prefer the return type to match the positive and negative arguments, which should be checked for match in size and shape. For 1D arrays, that would look like this

T[] if_else(int[] cond, T[] pos, T[] neg);

where

if_else(cond, pos, neg)[n]  =def=  cond[n] ? pos[n] : neg[n].

This could potentially be extended to higher-dimensional arrays, vectors, and matrices in the obvious way. For example,

T[ , ] if_else(int[ , ] cond, T[ , ] pos, T[ , ] neg;
vector[] if_else int[ , ] cond, vector[] pos, vector[] neg);
matrix if_else(int[ , ] cond, matrix pos, matrix neg);

The above definition is consistent with broadcasting a single cond, as in our current signature for the ternary op:

U operator?:=(int cond, U pos, U neg);

which applies to U = T[] to give

T[] operator?:=(int cond, T[] pos, T[] neg);

There's a question of whether if_else should be implemented for the scalar case because the ternary op is more efficient.

[bob-carpenter]

I'm working through mail backward and saw the note in the math lib about broadcasting. We could broadcast the arguments, but that can start to get confusing with the ternary operator. These two are natural consequences of having an unvectorized and fully vectorize form,

T[] if_else(int cond, T[] pos, T[] neg);  // ternary op
T[] if_else(int[] cond, T[] pos, T[] neg);  // fully vectorized

The natural broadcasting is this:

T[] if_else(int[] cond, T pos, T neg);
T[] if_else(int cond, T[] pos, T neg);
T[] if_else(int cond, T pos, T[] neg);
T[] if_else(int[] cond, T[] pos, T neg);
T[] if_else(int[] cond, T pos, T[] neg);

It feels dangerous given the ternary op and fully vectorized forms, but I don't see any cases where this'd run into ambiguity.

It gets trickier with matrices and vectors---we've always shied away from broadcasting vectors up to matrices by copying columsn or rwo vectors to matrices by copying rows.

[andrjohns]

Thanks Bob, that makes sense to me! In that case it could work to have stanc3 conditionally transpile to if_else() when any inputs are containers, otherwise default to the c++ ternary?

As for the broadcasting, the rules would be identical to the binary functions - all containers have to be the same dimension (i.e., no mixing vectors and matrices) and any scalars are recycled/broadcast - so the signatures that you posted above would be correct.

[bob-carpenter]

conditionally transpile to if_else() when any inputs are containers, otherwise default to the c++ ternary?

That's the only way I can see to maintain backward compatibility for short-circuiting.

all containers have to be the same dimension (i.e., no mixing vectors and matrices) and any scalars are recycled/broadcast

Perfect, thanks!

[WardBrian]

We end up doing something pretty similar for all the other operators - e.g. we use the building C++ operator if everything is a scalar, otherwise we call add():

https://github.com/stan-dev/stanc3/blob/9218def9e170e4d32ef0eecef8ec92b1dfbc6fdb/src/stan_math_backend/Expression_gen.ml#L218-L226