Tensor NestedTensor_slice()

in nestedtensor/csrc/nested_tensor_impl.cpp [253:310]

• 57 lines of code
• 13 McCabe index (conditional complexity)

Tensor NestedTensor_slice(
    const Tensor& self,
    int64_t dim,
    c10::optional<int64_t> start_,
    c10::optional<int64_t> end_,
    int64_t step) {
  int64_t start;
  if (start_) {
    start = *start_;
  } else {
    start = 0;
  }
  int64_t end;
  if (end_) {
    end = *end_;
  } else {
    end = 9223372036854775807;
  }
  int64_t ndim = get_dim(self);
  if (ndim == 0) {
    TORCH_CHECK_INDEX(false, "slice() cannot be applied to a 0-dim tensor.");
  }
  dim = maybe_wrap_dim(dim, ndim);
  if (dim != 0) {
    TORCH_CHECK_INDEX(false, "slice() only supports dim == 0 for now.");
  }
  // TODO: support negative strides
  TORCH_CHECK(step >= 1, "slice step must be positive for now.");
  int64_t sizes_0 = nt_size(self, 0);
  if (start < 0) {
    start += sizes_0;
  }
  if (end < 0) {
    end += sizes_0;
  }
  if (start < 0) {
    start = 0;
  } else if (start >= sizes_0) {
    start = sizes_0;
  }
  if (end < start) {
    end = start;
  } else if (end >= sizes_0) {
    end = sizes_0;
  }
  std::vector<at::Tensor> unbound = at::unbind(self, 0);
  std::vector<TensorNode> new_tensor_nodes;
  for (int64_t i = start; i < end; i += step) {
    if (is_nested_tensor_impl(unbound[i])) {
      new_tensor_nodes.push_back(get_nested_tensor_structure(unbound[i]));
    } else {
      new_tensor_nodes.push_back(TensorNode(std::move(unbound[i])));
    }
  }
  auto result = wrap_tensor_node(TensorNode(std::move(new_tensor_nodes)));
  namedinference::propagate_names(result, self);
  return result;
}