Fix TopK rank-1

py/torch_tensorrt/dynamo/conversion/impl/topk.py

@@ -194,26 +194,37 @@ def topk(
     sorted: Optional[bool],
     return_indices: bool = True,
 ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]:
-    if largest:
-        topk_layer = ctx.net.add_topk(
-            input,
-            trt.TopKOperation.MAX,
-            k,
-            get_axes_for_reduce_op(get_positive_dim(dim, len(input.shape))),
-        )
-    else:
-        topk_layer = ctx.net.add_topk(
-            input,
-            trt.TopKOperation.MIN,
-            k,
-            get_axes_for_reduce_op(get_positive_dim(dim, len(input.shape))),
+    # ITopKLayer requires rank >= 2 (same broadcast as argmax_argmin above).
+    is_rank_1 = len(input.shape) == 1
+    if is_rank_1:
+        input = impl.shuffle.reshape(
+            ctx, target, source_ir, f"{name}_broadcast", input, (*input.shape, 1)
         )
+        dim = 0
+
+    topk_layer = ctx.net.add_topk(
+        input,
+        trt.TopKOperation.MAX if largest else trt.TopKOperation.MIN,
+        k,
+        get_axes_for_reduce_op(get_positive_dim(dim, len(input.shape))),
+    )
 
     # TensorRT ITopKLayer does not have a sorted flag, it is always returning the sorted topk elements
     # so here no matter sorted is True or False the returned the topk Tensor object is always sorted
     set_layer_name(topk_layer, target, f"{name}_topk", source_ir)
 
+    values = topk_layer.get_output(0)
+    indices = topk_layer.get_output(1)
+
+    if is_rank_1:
+        values = impl.shuffle.reshape(
+            ctx, target, source_ir, f"{name}_values_1d", values, (-1,)
+        )
+        indices = impl.shuffle.reshape(
+            ctx, target, source_ir, f"{name}_indices_1d", indices, (-1,)
+        )
+
     if return_indices:
-        return topk_layer.get_output(0), topk_layer.get_output(1)
+        return values, indices
     else:
-        return topk_layer.get_output(0)
+        return values

tests/py/dynamo/conversion/test_sort_aten.py

@@ -17,6 +17,10 @@ class TestSortConverter(DispatchTestCase):
             ((1, 5, 2, 1), -1, True),
             ((1, 2, 5, 3), -2, False),
             ((6, 2, 1, 3), -4, True),
+            # ITopKLayer requires rank >= 2; converter broadcasts rank-1
+            ((64,), 0, False),
+            ((64,), 0, True),
+            ((16,), -1, True),
         ]
     )
     def test_sort(self, input_shape, dim, descending):

tests/py/dynamo/conversion/test_topk_aten.py

@@ -19,6 +19,10 @@ class TestSortConverter(DispatchTestCase):
             ((6, 4), 2, 1, False, False),
             # default dim:-1 largest:True, sorted:True
             ((3, 5, 12), 3),
+            # ITopKLayer requires rank >= 2; converter broadcasts rank-1
+            ((64,), 64, 0, True, True),
+            ((64,), 8, 0, False, True),
+            ((16,), 4, -1, True, True),
         ]
     )
     def test_topk(self, input_shape, k, dim=-1, largest=True, sorted=True):