[Code scan] Copy the innermost dimension when slicing 3D tensors

[njzjz]

This issue is a result of a Codex global repository scan.

Tensor::slice() copies 3D slices row by row. For each (i, j) row, offset_out advances by size[2], so the contiguous copy length should also be size[2]. The current code copies size[1], which corrupts non-cubic 3D slices where size[1] != size[2].

abacus-develop/source/source_base/module_container/ATen/core/tensor.cpp

Lines 191 to 200 in 84ca04b

	else if (ndim == 3) {
	for (int i = 0; i < size[0]; i++) {
	for (int j = 0; j < size[1]; j++) {
	int offset = static_cast<int>((i + start[0]) * shape_.dim_size(1) * shape_.dim_size(2) +
	(j + start[1]) * shape_.dim_size(2) + start[2]);
	int offset_out = i * size[1] * size[2] + j * size[2];
	TEMPLATE_ALL_2(this->data_type_, this->device_,
	kernels::synchronize_memory<T_, DEVICE_, DEVICE_>()(
	output.data<T_>() + offset_out, this->data<T_>() + offset, size[1]))
	}

Relevant code:

int offset_out = i * size[1] * size[2] + j * size[2];
TEMPLATE_ALL_2(this->data_type_, this->device_,
               kernels::synchronize_memory<T_, DEVICE_, DEVICE_>()(
                       output.data<T_>() + offset_out, this->data<T_>() + offset, size[1]))

Impact:

For size[1] < size[2], the tail of each row is left uninitialized. For size[1] > size[2], the copy writes past the intended output row and can corrupt the next row.

Suggested fix:

Change the 3D copy length to size[2] and add a regression test for a non-cubic 3D slice, for example size = {2, 3, 5}.