Adds shape propagation to Q/DQ ops

Author: pranavm-nvidia

tripy/nvtripy/trace/ops/cast.py

@@ -26,9 +26,7 @@
 class Cast(TraceOp):
     dtype: "nvtripy.common.dtype"
 
-    def infer_rank(self):
-        # Casting does not change the shape, so we can simply copy it.
-        self.outputs[0].shape = self.inputs[0].shape
+    infer_rank = op_utils.InferRankPolicies.same_shape_as_input()
 
     def infer_dtypes(self):
         self.outputs[0].dtype = self.dtype

tripy/nvtripy/trace/ops/dequantize.py

@@ -29,7 +29,7 @@ class Dequantize(TraceOp):
     dtype: datatype.dtype
     dim: int
 
-    infer_rank = op_utils.InferRankPolicies.same_as_input()
+    infer_rank = op_utils.InferRankPolicies.same_shape_as_input()
 
     def infer_dtypes(self):
         self.outputs[0].dtype = self.dtype

tripy/nvtripy/trace/ops/quantize.py

@@ -29,7 +29,7 @@ class Quantize(TraceOp):
     dtype: datatype.dtype
     dim: int
 
-    infer_rank = op_utils.InferRankPolicies.same_as_input()
+    infer_rank = op_utils.InferRankPolicies.same_shape_as_input()
 
     def infer_dtypes(self):
         self.outputs[0].dtype = self.dtype

tripy/nvtripy/trace/ops/utils.py

@@ -23,6 +23,13 @@
 
 
 class InferRankPolicies:
+    # Indicates that the output has not only the same rank but also the same shape as the input.
+    def same_shape_as_input(idx=0):
+        def impl(self):
+            self.outputs[0].shape = self.inputs[idx].shape
+
+        return impl
+
     def same_as_input(idx=0):
         def impl(self):
             self.outputs[0].rank = self.inputs[idx].rank