apply_op Class — pytorch Architecture

Source Code

aten/src/ATen/CPUApplyUtils.h lines 234–256

template <typename Op, typename... Args>
inline void apply_op(
    int64_t numel,
    int64_t offset,
    const Op& op,
    Args... iters) {
  // For 0-dim tensors
  if (numel == 1 && max_dim(iters...) == 0) {
    op(*iters.data_...);
    return;
  }
  if (offset > 0)
    forward(offset, iters...);
  // Splitting this into chunks helps the compiler create faster assembly
  for (int64_t i = 0; i < numel;) {
    for (; iterate_continue(iters...) && i < numel;) {
      op(*iters.data_...);
      iterate(1, iters...);
      i++;
    }
    iterate_overflow(iters...);
  }
}