Skip to content

Commit

Permalink
Fix linalg iterator. (#8603)
Browse files Browse the repository at this point in the history
  • Loading branch information
trivialfis authored Dec 16, 2022
1 parent 38887a1 commit a10e4cb
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions src/common/linalg_op.h
Original file line number Diff line number Diff line change
Expand Up @@ -63,27 +63,27 @@ void ElementWiseKernel(Context const* ctx, linalg::TensorView<T, D> t, Fn&& fn)
#endif // !defined(XGBOOST_USE_CUDA)

template <typename T, std::int32_t kDim>
auto cbegin(TensorView<T, kDim> v) { // NOLINT
auto cbegin(TensorView<T, kDim> const& v) { // NOLINT
auto it = common::MakeIndexTransformIter([&](size_t i) -> std::remove_cv_t<T> const& {
return linalg::detail::Apply(v, linalg::UnravelIndex(i, v.Shape()));
});
return it;
}

template <typename T, std::int32_t kDim>
auto cend(TensorView<T, kDim> v) { // NOLINT
auto cend(TensorView<T, kDim> const& v) { // NOLINT
return cbegin(v) + v.Size();
}

template <typename T, std::int32_t kDim>
auto begin(TensorView<T, kDim> v) { // NOLINT
auto begin(TensorView<T, kDim>& v) { // NOLINT
auto it = common::MakeIndexTransformIter(
[&](size_t i) -> T& { return linalg::detail::Apply(v, linalg::UnravelIndex(i, v.Shape())); });
return it;
}

template <typename T, std::int32_t kDim>
auto end(TensorView<T, kDim> v) { // NOLINT
auto end(TensorView<T, kDim>& v) { // NOLINT
return begin(v) + v.Size();
}
} // namespace linalg
Expand Down

0 comments on commit a10e4cb

Please sign in to comment.