* Added offset check in Relay.

* Minor changes.
apache · Sep 23, 2020 · b44a394 · b44a394
1 parent 8bff88c
commit b44a394
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
@@ -1541,7 +1541,7 @@ inline Tensor matrix_set_diag(const Tensor& input, const Tensor& diagonal, int k
                               const std::string tag = kInjective) {
   size_t ndim = input->shape.size() - 1;
 
-  bool onlyOneDiagonal = k1 == k2;
+  bool only_one_diagonal = k1 == k2;
 
   return compute(
       input->shape,
@@ -1552,7 +1552,7 @@ inline Tensor matrix_set_diag(const Tensor& input, const Tensor& diagonal, int k
           for (size_t i = 0; i < ndim - 1; i++) {
             diagonal_indices.push_back(iter_vars[i]);
           }
-          if (onlyOneDiagonal) {
+          if (only_one_diagonal) {
             k = k1;
           } else {
             // Determining which diagonal/sub-diagonal/super-diagonal it is

diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc
@@ -3138,6 +3138,10 @@ bool MatrixSetDiagRel(const Array<Type>& types, int num_inputs, const Attrs& att
 
   int d_ndims = diagonal->shape.size();
   int i_ndims = input->shape.size();
+
+  reporter->Assert(input->shape[i_ndims - 2] > -param->k1);
+  reporter->Assert(input->shape[i_ndims - 1] > param->k2);
+
   for (int i = 0; i < d_ndims - 2; i++) {
     reporter->AssertEQ(input->shape[i], diagonal->shape[i]);
   }