stan-dev / math

The Stan Math Library is a C++ template library for automatic differentiation of any order using forward, reverse, and mixed modes. It includes a range of built-in functions for probabilistic modeling, linear algebra, and equation solving.
https://mc-stan.org
BSD 3-Clause "New" or "Revised" License
756 stars 188 forks source link

allow `array[] int` to add and subtract elementwise an `int` or `array[] int` of the same size #2782

Open spinkney opened 2 years ago

spinkney commented 2 years ago

With vectorized indexing I expected to be able to increment/decrement the array[] int but that is not allowed. For this special case I think it makes sense to allow adding and subtracting elementwise into Stan.

bob-carpenter commented 2 years ago

@spinkney: Could you provide an example and ideally the signatures of operators or functions you'd like to add? Given that we do allow increments and decrements of array elements, I'm guessing you're talking about adding two integer arrays together or maybe adding a scalar to an int array? So I'm guessing you mean these three signatures or maybe just the first.

array[] int add(array[] int x, array[] int y);

array[] int add(int x, array[] int y);
array[] int add(array[] int x, int y);

Then presumably you also want the subtract forms? What about other elementwise operations like multiplication and division?

I do not want to allow mixed operations with arrays of reals---you can use to_vector on an integer array and matrix arithmetic when the result is a vector or matrix.

The goal in writing an issue is enough detail that someone can implement it without guesswork. The problem with vague issues is that someone might take a guess at how to implement then get their PR rejected because it wasn't what the author of the issue intended. Usually we encourage discussion on the forum if a discussion is needed to clarify an issue before posting it.

spinkney commented 2 years ago

Yes, I'd like those 3 cases for addition and subtraction. For example, I had some code like below

    int length_x;
    array[length_x] int x = linspaced_int_array(length_x, L, U);
    int k;
    vector[k + 1] lsum;

    for(i in 1:length_x) {
      t[i] = -lsum[x[i] + 1] - lsum[mA - x[i] + 1] - lsum[n - x[i] + 1] - lsum[mB - n + x[i] + 1];
    }

I initially didn't have the loop thinking that the x + 1 would work but since x is an array of ints it errors out.

I don't know about the multiplication and division cases. I guess those could come in handy as well but we'd need to round down to int if needed.

rok-cesnovar commented 2 years ago

I think in this case, the easiest thing to do is just create utility UDFs and use them. With shadowing and overloading this is now pretty simple:

functions {
    array[] int add(array[] int x, array[] int y) {
        int x_size = size(x);
        array[x_size] int z;
        for (i in 1:x_size){
          z[i] = x[i] + y[i];
        }
        return z;
    }
    array[] int add(int x, array[] int y) {
        int y_size = size(y);
        array[y_size] int z;
        for (i in 1:y_size){
          z[i] = x + y[i];
        }
        return z;
    }
    array[] int add(array[] int x, int y) {
        int x_size = size(x);
        array[x_size] int z;
        for (i in 1:x_size){
          z[i] = x[i] + y;
        }
        return z;
    }
}
transformed data {
    array[5] int x;
    array[5] int y;
    array[5] int z1 = add(x, y);
    array[5] int z2 = add(x, 5);
    array[5] int z3 = add(3, y);
}

Given that add() is just the function call for the +operator this should probably work with the + as well, but it does not. I think if we were able to write the following with the above UDFs, that would be good enough.

    array[5] int z1 = x + y;
    array[5] int z2 = x + 5;
    array[5] int z3 = 3 + y;

There is no efficiency to be gained here with a Stan Math implementation, so not sure this is something worth a separate Math implementation. I can always be persuaded otherwise :) Incrementing arrays of indices might be something a lot of users do?

WardBrian commented 2 years ago

I believe I suggested to the contrary on slack recently, but I should note that the user defining functions called things like add, subtract, etc right now does not allow those to be used with +, -, etc. There's no hard technical constraint here, the typechecker just looks up specifically library functions when doing that typecheck.

spinkney commented 2 years ago

I realize there's no efficiency gain. As a user, I just expected the vectorized indexing to work with addition/subtraction and array integers.

bob-carpenter commented 2 years ago

The only reason I'm OK with this for int is that we don't have the vector alternative. Most of our users coming from Python and R expect arithmetic to work on real arrays, too. My question is how much arithmetic do we need to implement so it's not confusing. Presumably we don't implement addition without also implementing subtraction and negation. We could also included elementwise multiplication or division.

rok-cesnovar commented 2 years ago

If we only do this for array[] int for basic arithmetic, that would be fine. Also not a ton of work.

andrjohns commented 2 years ago

If we add a specialisation using apply_scalar_binary:

template <typename T1, typename T2,
          require_any_std_vector_t<T1, T2>* = nullptr>
inline auto add(const T1& a, const T2& b) {
  return apply_scalar_binary(
      a, b, [](const auto& c, const auto& d) { return add(c, d); });
}

This will handle all arrays/nested containers