[Code scan] Validate neighbor indices and mm tensor rank in edge_index ops

[njzjz]

This issue was found by a Codex global scan of the repository at commit 19f9265.

The CPU and CUDA edge builders treat every kk >= 0 neighbor as valid, but they do not check kk < nall before using it to form global_kk and read atype[global_kk].

deepmd-gnn/op/edge_index.cc

Lines 90 to 98 in 19f9265

	const int64_t kk = nlist[idx];
	const int64_t safe_kk = std::max<int64_t>(kk, 0);
	const int64_t global_kk = ff * shape.nall + safe_kk;
	const int64_t global_ii = ff * shape.nall + ii;
	bool valid = kk >= 0;
	if (valid && nmm > 0) {
	const bool in_mm1 = type_in_mm(atype[global_ii], mm, nmm);
	const bool in_mm2 = type_in_mm(atype[global_kk], mm, nmm);
	valid = !(in_mm1 && in_mm2);

deepmd-gnn/op/edge_index_cuda.cu

Lines 52 to 70 in 19f9265

	const int64_t kk = nlist[tid];
	if (kk < 0) {
	return false;
	}
	const int64_t global_ii = ff * nall + ii;
	const int64_t global_kk = ff * nall + kk;
	bool in_mm1 = false;
	for (int64_t mm_idx = 0; mm_idx < nmm; ++mm_idx) {
	if (atype[global_ii] == mm[mm_idx]) {
	in_mm1 = true;
	break;
	}
	}
	bool in_mm2 = false;
	for (int64_t mm_idx = 0; mm_idx < nmm; ++mm_idx) {
	if (atype[global_kk] == mm[mm_idx]) {

The shape validation also does not check that nloc <= nall, and the mm tensor rank is not validated before using size(0) as the count:

deepmd-gnn/op/edge_index.cc

Lines 32 to 58 in 19f9265

	EdgeIndexShape validate_shape(const ts::Tensor& nlist_tensor,
	const ts::Tensor& atype_tensor) {
	EdgeIndexShape shape;
	if (nlist_tensor.dim() == 2) {
	if (atype_tensor.dim() != 1) {
	throw std::invalid_argument("atype_tensor must be 1D");
	}
	shape.nf = 1;
	shape.nloc = nlist_tensor.size(0);
	shape.nnei = nlist_tensor.size(1);
	shape.nall = atype_tensor.size(0);
	} else if (nlist_tensor.dim() == 3) {
	if (atype_tensor.dim() != 2) {
	throw std::invalid_argument("atype_tensor must be 2D");
	}
	shape.nf = nlist_tensor.size(0);
	shape.nloc = nlist_tensor.size(1);
	shape.nnei = nlist_tensor.size(2);
	if (atype_tensor.size(0) != shape.nf) {
	throw std::invalid_argument(
	"atype_tensor must have the same size as nlist_tensor");
	}
	shape.nall = atype_tensor.size(1);
	} else {
	throw std::invalid_argument("nlist_tensor must be 2D or 3D");
	}
	return shape;

deepmd-gnn/op/edge_index.cc

Lines 113 to 117 in 19f9265

	const EdgeIndexShape shape = validate_shape(nlist_tensor_, atype_tensor_);
	const int64_t nmm = mm_tensor_.size(0);
	const int64_t* nlist = nlist_tensor_.const_data_ptr<int64_t>();
	const int64_t* atype = atype_tensor_.const_data_ptr<int64_t>();
	const int64_t* mm = mm_tensor_.const_data_ptr<int64_t>();

deepmd-gnn/op/edge_index.cc

Lines 184 to 188 in 19f9265

	const EdgeIndexShape shape = validate_shape(nlist_tensor_, atype_tensor_);
	const int64_t nmm = mm_tensor_.size(0);
	const int64_t* nlist = nlist_tensor_.const_data_ptr<int64_t>();
	const int64_t* atype = atype_tensor_.const_data_ptr<int64_t>();
	const int64_t* mm = mm_tensor_.const_data_ptr<int64_t>();

Bad direct torch.ops.deepmd_gnn.edge_index(...) inputs can therefore produce out-of-bounds CPU/GPU reads or invalid edge indices. This is narrower than the already-closed mapping-path issue #58: this one is about validating the native op's own inputs.

Suggested fix: reject atype shapes where nall < nloc, require mm_tensor.dim() == 1, and either throw on non-negative neighbor indices >= nall or mark them invalid consistently in both sparse and dense paths.