Fix style issue and respond to feedback.

python/tvm/topi/nn/batch_matmul.py

@@ -42,7 +42,7 @@ def batch_matmul(x, y):
     y_shape = get_const_tuple(y.shape)
     XB = x_shape[0]
     YB = y_shape[0]
-    assert (XB == YB) or (XB == 1) or (YB == 1), "batch dimension doesn't match"
+    assert (XB == YB) or (XB == 1) or (YB == 1), "batch dimensions don't match"
     assert x_shape[2] == y_shape[2], "shapes of x and y is inconsistant"
     _, M, K = x.shape
     batch = max(XB, YB)

src/relay/op/nn/nn.cc

@@ -35,6 +35,7 @@
 
 #include <string>
 #include <vector>
+#include <algorithm>
 
 #include "../../transforms/infer_layout_util.h"
 #include "../make_op.h"
@@ -851,10 +852,10 @@ bool BatchMatmulRel(const Array<Type>& types, int num_inputs, const Attrs& attrs
   const auto* y = types[1].as<TensorTypeNode>();
   if (x == nullptr || y == nullptr) return false;
   CHECK(x->shape.size() == 3 && y->shape.size() == 3);
-  CHECK(reporter->AssertEQ(x->shape[0], y->shape[0]) || 
+  CHECK(reporter->AssertEQ(x->shape[0], y->shape[0]) ||
         reporter->AssertEQ(x->shape[0], 1) ||
         reporter->AssertEQ(y->shape[0], 1))
-      << "BatchDot: batch dimension doesn't match, "
+      << "BatchDot: batch dimensions don't match, "
       << " x shape=" << x->shape << ", y shape=" << y->shape;
   CHECK(reporter->AssertEQ(x->shape[2], y->shape[2]))
       << "BatchDot: shapes of x and y is inconsistent, "

tests/python/topi/python/test_topi_dense.py

@@ -75,8 +75,6 @@ def check_device(device, ctx):
             b = tvm.nd.array(b_np, ctx)
             c = tvm.nd.array(c_np, ctx)
             d = tvm.nd.array(np.zeros(get_const_tuple(D.shape), dtype=dtype), ctx)
-            print(tvm.lower(s, [A, B, C, D], simple_mode=True))
-            exit()
             f = tvm.build(s, [A, B, C, D], device, name="dense")
             f(a, b, c, d)
             tvm.testing.assert_allclose(d.asnumpy(), d_np, rtol=1e-5)