Elementwise functions & Eigen::Map

Description

Now that the apply_vector_unary framework has gone through, I want to use it to extend the elementwise functions(e.g. exp, acos). At the moment, Eigen's SIMD-optimised functions are only used for Eigen::MatrixBase objects (e.g. Eigen::Matrix and Eigen::Map), not for std::vectors (containing either scalars or Eigen objects) or for Eigen::Array objects (unless an additional specialisation is added, as was done for some functions). This leaves a bit of speed on the table, as these objects then get called with the less optimised versions of the function.

Extending these functions to be agnostic betweenstd::vector and Eigen::Array is simple with the apply_vector_unary approach. The difficulty comes in the return type. Because the rest of the Stan library is not yet able to take Eigen expressions as inputs, the function needs to be .eval-ed to return either an Eigen::Matrix or an Eigen::Array, depending on the input.

From what I can see, there are two ways of achieving this. We could have a different specialisation for MatrixBase and ArrayBase, however this means the entire function would be duplicated, with the only difference being the return.

Using exp as an example:

template <typename Derived,
          typename = require_eigen_vt<std::is_arithmetic, Derived>>
inline auto exp(const Eigen::MatrixBase<Derived>& x) {
  return x.derived().array().exp().matrix().eval();
}

template <typename Derived,
          typename = require_eigen_vt<std::is_arithmetic, Derived>>
inline auto exp(const Eigen::ArrayBase<Derived>& x) {
  return x.derived().array().exp().eval();
}

Alternatively, we could use a helper function that uses the input type to determine the right wrapper to call (.matrix() or .array()) prior to the call to .eval(). As an initial prototype, I've created the function match_wrapper<T>() which will call the .array() wrapper for any Eigen::Array objects, and the .matrix() wrapper otherwise.

Using exp as an example, the following function will use Eigen's exp function for all containers (and containers of containers) and return either an Eigen::Matrix or Eigen::Array (or container of these), depending on the input type:

template <typename T, require_container_st<is_container, std::is_arithmetic, T>>
inline auto exp(const T& x) {
  return apply_vector_unary<T>::apply(x, [&](const auto& v) {
    return match_wrapper<decltype(v)>(v.derived().array().exp()).eval();
  });
}

Where match_wrapper is defined as:

template <typename T1, typename T2, require_t<is_eigen_array<T1>>...>
const auto& match_wrapper(const T2& x) { return x.array(); }

template <typename T1, typename T2, require_not_t<is_eigen_array<T1>>...>
const auto match_wrapper(const T2& x) { return x.matrix(); }

Does this sound reasonable? Any concerns/suggestions about the approach?

Current Version:

v3.1.0

Alternatively, we could use a helper function that uses the input type to determine the right wrapper to call (.matrix() or .array()) prior to the call to .eval(). As an initial prototype, I've created the function match_wrapper() which will call the .array() wrapper for any Eigen::Array objects, and the .matrix() wrapper otherwise.

Does this work? I thought Eigen::Matrix does not define .matrix() function and Eigen::Array does not define .array().

Anyway, I like this idea. I prefere wrapper function, but I think it could be used in apply_vector_unary instead of everywhere apply_vector_unary is used.

Does this work? I thought Eigen::Matrix does not define .matrix() function and Eigen::Array does not define .array().

Haven't had any problems so far (touch wood), I've got a branch here testing the replacements, and the unit tests all pass without issue.

I think it could be used in apply_vector_unary instead of everywhere apply_vector_unary is used.

I'm not sure this would be preferable, since doesn't the wrapper need to be called before the .eval()? That would require moving the eval out to apply_vector_unary as well, which would introduce the extra copy for std::vectors that you mentioned when you fixed up the expression returns.

I don't see any new tests on your branch. What I predict would would fail with your code, is calling, for example exp(Eigen::Array), which was not previously suported and therefore not in tests.

I think we don't call .eval() in apply_vector_unary anyway. I see this on your branch:

  return apply_vector_unary<T>::apply(x, [&](const auto& v) {
    return match_wrapper<decltype(v)>(v.derived().array().asin()).eval();
  });

which could be simplified to

  return apply_vector_unary<T>::apply(x, [&](const auto& v) {
    return v.derived().array().asin();
  }).eval();

and calling match_wrapper in apply_vector_unary. Or you could include even .derived().array() in apply_vector_unary (.array() would have to be replaced with another function similar to match_wrapper).

Also another question: Do we actually want to introduce support for calling element-wise functions, such as exp on std::vectors? I don' t have any preference on this, but I got a feeling some people might be against this.

EDIT: fixed suggested change

I don't see any new tests on your branch. What I predict would would fail with your code, is calling, for example exp(Eigen::Array), which was not previously suported and therefore not in tests.

Works like a charm, I've pushed an updated test to the branch so you can verify.

Ah, I see what you mean about moving the wrapper, looks pretty good to me, I'll do some testing.

Also another question: Do we actually want to introduce support for calling element-wise functions, such as exp on std::vectors? I don' t have any preference on this, but I got a feeling some people might be against this.

It's actually already supported (but not tested) through apply_scalar_unary, so this is just optimising existing functionality

Edit: It's tested in mix, but not in prim

If the argument is a std::vector, the return type needs to be std::vector.

If the argument is a std::vector, the return type needs to be std::vector.

Yep, std::vector inputs will still return std::vector outputs, Eigen::Map is just applied so that Eigen's exp implementation can be used.

The Matrix/Array return type discussion above is just working out a way to handle both Eigen::Matrix and Eigen::Array inputs, while returning the appropriate type, without duplicating the function.

@t4c1

@t4c1 When you've got a minute, could I get hand with a segfault that I can't figure out? I'm attempting to move the match_wrapper call to apply_vector_unary, but it's causing a segfault in the mix tests for log_softmax (which uses the framework). Oddly enough, the segfault doesn't occur when I explicitly call .matrix() or .array() on the return, just when I use match_wrapper.

This makes me think that it's got something to do with the passing/returning to or from match_wrapper, but I can't figure out where the issue is. I've got a branch up here which has the minimal changes needed to cause a segfault (but only in the mix tests for log_softmax).

What do you think?

- return (v.array() - log_sum_exp(v)).matrix().eval(); + return (v.array() - log_sum_exp(v)); The template expression for the subtraction holds references to local variables. The eval() forces the return type to be a copy.

…

On Feb 18, 2020, at 8:43 PM, Andrew Johnson ***@***.***> wrote: @t4c1 When you've got a minute, could I get hand with a segfault that I can't figure out? I'm attempting to move the match_wrapper call to apply_vector_unary, but it's causing a segfault in the mix tests for log_softmax (which uses the framework). Oddly enough, the segfault doesn't occur when I explicitly call .matrix() or .array() on the return, just when I use match_wrapper. This makes me think that it's got something to do with the passing/returning to or from match_wrapper, but I can't figure out where the issue is. I've got a branch up here which has the minimal changes needed to cause a segfault (but only in the mix tests for log_softmax). What do you think? — You are receiving this because you commented. Reply to this email directly, view it on GitHub, or unsubscribe.

That's a good point, and I would think that would be the issue, but it works fine depending on how the wrapper is called in apply_vector_unary.

This works:

  template <typename F>
  static inline auto apply(const T& x, const F& f) {
    return f(x).matrix().eval();
  }

But this doesn't:

  template <typename F>
  static inline auto apply(const T& x, const F& f) {
    return match_wrapper<T>(f(x)).eval();
  }

Where match_wrapper is defined as:

template <typename T1, typename T2>
const auto match_wrapper(const T2& x) { return x.matrix(); }

Further, adding .eval() to log_softmax itself also doesn't resolve the issue:

template <typename T, require_t<std::is_arithmetic<scalar_type_t<T>>>...>
inline auto log_softmax(const T& x) {
  return apply_vector_unary<T>::apply(x, [&](const auto& v) {
    check_nonzero_size("log_softmax", "v", v);
    return (v.array() - log_sum_exp(v)).eval();
  });
}

It looks like the way forward is to avoid match_wrapper entirely and instead template different versions of apply_vector_unary::apply for Eigen inputs, one for returning matrices and one for returning arrays:

// Return Eigen::Matrix
  template <typename F, typename T2 = T,
            require_t<is_eigen_matrix<plain_type_t<T2>>>...>
  static inline auto apply(const T& x, const F& f) {
    return f(x).matrix().eval();
  }

// Return Eigen::Array
  template <typename F, typename T2 = T,
            require_t<is_eigen_array<plain_type_t<T2>>>...>
  static inline auto apply(const T& x, const F& f) {
    return f(x).array().eval();
  }

Then for std::vector inputs, we know that the .matrix() wrapper should be used, since they'll be mapped to column vectors, and we can assign to them directly without first .eval()-ing the expression into a temporary:

  template <typename F>
  static inline std::vector<T_vt> apply(const T& x, const F& f) {
    std::vector<T_vt> result(x.size());
    Eigen::Map<Eigen::Matrix<T_vt, -1, 1>>(result.data(), result.size())
        = f(as_column_vector_or_scalar(x)).matrix();
    return result;
  }

Also, by having the single .eval() here rather than in each of the functions, it will be much easier to change once the rest of the library can accept expressions as inputs.

I see nothing obviously wrong with your code. After some debugging I found that rev version with nonempty input segfaults when returning from match_wrapper. Than I got stuck.

I figured that it works if you skip .matrix() call in match_wrapper (which is not necessary for matrices anyway). However, I don't know why that change makes it work.

Closed via #1727

Navigation Menu

Repository navigation

Elementwise functions & Eigen::Map #1723

Description

Current Version:

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Elementwise functions & Eigen::Map #1723

Description

Description

Current Version:

Activity

t4c1 commented on Feb 18, 2020

andrjohns commented on Feb 18, 2020

t4c1 commented on Feb 18, 2020

andrjohns commented on Feb 18, 2020

bob-carpenter commented on Feb 18, 2020

andrjohns commented on Feb 18, 2020

andrjohns commented on Feb 19, 2020

bob-carpenter commented on Feb 19, 2020

andrjohns commented on Feb 19, 2020

andrjohns commented on Feb 19, 2020

t4c1 commented on Feb 19, 2020

andrjohns commented on Mar 18, 2020

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions