Fix gather_nd in Relay (#3442)

* Fix gather_nd in Relay

* Add test cases for gather_nd.

src/relay/op/tensor/transform.cc

@@ -2332,7 +2332,7 @@ bool GatherNDRel(const Array<Type>& types,
     return false;
   }
   const size_t ndim = data->shape.size();
-  const IntImm* mdim = data->shape[0].as<IntImm>();
+  const IntImm* mdim = indices->shape[0].as<IntImm>();
   const size_t kdim = indices->shape.size() - 1;
   CHECK(size_t(mdim->value) <= ndim)
         << "GatherND: indices shape does satisfy.";

tests/python/frontend/mxnet/test_forward.py

@@ -528,6 +528,8 @@ def test_forward_gather_nd():
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res.asnumpy())
     verify((2, 2), (2, 3), [[1, 1, 0], [0, 1, 0]])
     verify((2, 2, 2), (2, 2), [[0, 1], [1, 0]])
+    verify((3, 2, 2), (2, 2), [[0, 1], [1, 0]])
+    verify((3, 2), (2, 2, 3), [[[0, 1, 2], [2, 0, 1]], [[0, 0, 0], [1, 1, 1]]])
 
 def test_forward_bilinear_resize():
     # add tests including scale_height and scale_width when mxnet is updated to version 1.5

tests/python/relay/test_op_level3.py

@@ -682,7 +682,8 @@ def test_gather_nd():
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
     verify_gather_nd((2, 2), (2, 3), [[1, 1, 0], [0, 1, 0]])
     verify_gather_nd((2, 2, 2), (2, 2), [[0, 1], [1, 0]])
-
+    verify_gather_nd((3, 2, 2), (2, 2), [[0, 1], [1, 0]])
+    verify_gather_nd((3, 2), (2, 2, 3), [[[0, 1, 2], [2, 0, 1]], [[0, 0, 0], [1, 1, 1]]])
 
 if __name__ == "__main__":
     test_arange()