index_add GPU op check rank of indices tensor

This commit is contained in:
tom__chen 2021-03-29 11:28:17 -04:00
parent b802f25563
commit ea02bf6e7a
1 changed files with 1 additions and 0 deletions

View File

@ -4601,6 +4601,7 @@ class IndexAdd(PrimitiveWithInfer):
validator.check("x rank", len(x_shape), "y rank", len(y_shape), Rel.EQ, self.name)
x_rank = len(x_shape)
validator.check_int_range(self.axis, -x_rank - 1, x_rank, Rel.INC_NEITHER, 'axis', self.name)
validator.check_equal_int(len(idx_shape), 1, "rank of idx_shape", self.name)
validator.check("size of indices", idx_shape[0], "dimension of y[axis]", y_shape[self.axis],
Rel.EQ, self.name)
axis = self.axis if self.axis >= 0 else x_rank + self.axis