chore: updates

Author: peri044

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

@@ -452,6 +452,7 @@ def pdist(
     p: float = 2,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     shape = input.shape
+    # Extend input from shape [N, D] to [N, 1, D]
     extend_input = impl.shuffle.reshape(
         ctx,
         target,
@@ -460,7 +461,18 @@ def pdist(
         input,
         shape=shape[0:1] + (1,) + shape[1:],
     )
-    x = impl.elementwise.sub(ctx, target, source_ir, f"{name}_sub", extend_input, input)
+    # Expand the input from [N, 1, D] to [N, N, D]
+    x = impl.slice.expand(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_sub",
+        extend_input,
+        (shape[0], shape[0]) + shape[1:],
+    )
+    # Subtract the expanded input from original input. Result shape = [N, N, D]
+    # This matrix has the distance of each sample to every other sample and hence the shape is [N, N, D]
+    x = impl.elementwise.sub(ctx, target, source_ir, f"{name}_sub", x, input)
 
     if p == 0:
         # norm = torch.sum(x!=0, dim=2)

tests/py/dynamo/conversion/test_layer_norm_aten.py

@@ -24,31 +24,6 @@ def forward(self, x):
             inputs,
         )
 
-    def test_layernorm_with_dynamic_shape(self):
-        class LayerNorm(torch.nn.Module):
-            def forward(self, x):
-                return torch.ops.aten.layer_norm.default(
-                    x,
-                    torch.tensor([3, 224, 224]),
-                    torch.ones((3, 224, 224)),
-                    torch.zeros((3, 224, 224)),
-                    1e-05,
-                    True,
-                )
-
-        input_specs = [
-            Input(
-                shape=(-1, 3, 224, 224),
-                dtype=torch.float32,
-                shape_ranges=[((1, 3, 224, 224), (1, 3, 224, 224), (2, 3, 224, 224))],
-            ),
-        ]
-
-        self.run_test_with_dynamic_shape(
-            LayerNorm(),
-            input_specs,
-        )
-
 
 class TestNativeLayerNormConverter(DispatchTestCase):
     def test_layer_norm(self):
@@ -68,30 +43,6 @@ def forward(self, x):
             inputs,
         )
 
-    def test_layernorm_with_dynamic_shape(self):
-        class LayerNorm(torch.nn.Module):
-            def forward(self, x):
-                return torch.ops.aten.native_layer_norm.default(
-                    x,
-                    torch.tensor([3, 224, 224]),
-                    torch.ones((3, 224, 224)),
-                    torch.zeros((3, 224, 224)),
-                    1e-05,
-                )[0]
-
-        input_specs = [
-            Input(
-                shape=(-1, 3, 224, 224),
-                dtype=torch.float32,
-                shape_ranges=[((1, 3, 224, 224), (1, 3, 224, 224), (2, 3, 224, 224))],
-            ),
-        ]
-
-        self.run_test_with_dynamic_shape(
-            LayerNorm(),
-            input_specs,
-        )
-
 
 if __name__ == "__main__":
     run_tests()