From b44a394e36b788029ad72a5b6daa212e5aec6002 Mon Sep 17 00:00:00 2001 From: Rishabh Jain Date: Wed, 23 Sep 2020 16:08:41 +0100 Subject: [PATCH] * Added offset check in Relay. * Minor changes. --- include/tvm/topi/transform.h | 4 ++-- src/relay/op/tensor/transform.cc | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h index 0cc97c5bb039..e01eb703cb99 100644 --- 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 index 9deabb597f41..4faface03ac6 100644 --- a/src/relay/op/tensor/transform.cc +++ b/src/relay/op/tensor/transform.cc @@ -3138,6 +3138,10 @@ bool MatrixSetDiagRel(const Array& 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]); }