slice Class — pytorch Architecture

Source Code

aten/src/ATen/native/transformers/attention.cpp lines 617–629

template <bool slice>
at::Tensor pad_last_dim(const at::Tensor& attn_bias, int alignment_size) {
  auto last_dim_size = attn_bias.sym_size(-1);
  if (last_dim_size % alignment_size == 0) {
    return attn_bias;
  }
  auto pad_count = alignment_size - (last_dim_size % alignment_size);
  auto padded_bias = at::pad_symint(attn_bias, {c10::SymInt(0), pad_count});
  if (slice) {
    return padded_bias.slice_symint(-1, 0, last_dim_size);
  }
  return padded_bias;
}