arange Class — pytorch Architecture

Source Code

aten/src/ATen/cpu/vec/vec128/vec128_uint_aarch64.h lines 282–291

template <typename step_t>
inline Vectorized<uint8_t> Vectorized<uint8_t>::arange(
    uint8_t base,
    step_t step) {
  const Vectorized<uint8_t> base_vec(base);
  const Vectorized<uint8_t> step_vec(step);
  const uint8x16_t step_sizes = {
      0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
  return vmlaq_u8(base_vec, step_sizes, step_vec);
}