[GitHub] [tvm] trevor-m commented on a change in pull request #7026: [BYOC][TRT] Support batch norm for all ranks <=5, and all axes

[GitBox <gi...@apache.org>]

trevor-m commented on a change in pull request #7026:
URL: https://github.com/apache/tvm/pull/7026#discussion_r537729585



##########
File path: src/runtime/contrib/tensorrt/tensorrt_ops.cc
##########
@@ -386,8 +386,35 @@ class BatchNormOpConverter : public TensorRTOpConverter {
     const int axis = std::stoi(params->node.GetAttr<std::vector<std::string>>("axis")[0]);
     const bool scale = std::stoi(params->node.GetAttr<std::vector<std::string>>("scale")[0]);
     const bool center = std::stoi(params->node.GetAttr<std::vector<std::string>>("center")[0]);
-    ICHECK(axis == 1 || axis == 3);
-    const bool need_transpose = axis == 3;
+    auto input_dims = TrtDimsToVector(input->getDimensions());
+    const size_t min_rank = TRT_HAS_IMPLICIT_BATCH(params) ? 3 : 4;
+    const size_t max_rank = TRT_HAS_IMPLICIT_BATCH(params) ? 4 : 5;
+    ICHECK_LE(input_dims.size(), max_rank);

Review comment:
       Hi @jroesch, thanks for reviewing!
   
   These checks are more for sanity checking, since the annotation functions in python/tvm/relay/op/contrib/tensorrt.py will filter out the unsupported ops before they ever get to this code. I don't expect users to ever see these.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org