resolve comments

Author: bowang007

.github/workflows/build-test-linux.yml

@@ -143,6 +143,7 @@ jobs:
         python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml -n 4 conversion/
         python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml automatic_plugin/test_automatic_plugin.py
         python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml automatic_plugin/test_automatic_plugin_with_attrs.py
+        python -m pytest -ra --junitxml=${RUNNER_TEST_RESULTS_DIR}/dynamo_converters_test_results.xml automatic_plugin/flashinfer_plugin.py
         popd
 
   tests-py-dynamo-fe:

py/requirements.txt

@@ -6,3 +6,4 @@ torch>=2.7.0.dev,<2.8.0
 torchvision>=0.22.0.dev,<0.23.0
 --extra-index-url https://pypi.ngc.nvidia.com
 pyyaml
+flashinfer-python

py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin.py

@@ -1,3 +1,4 @@
+import itertools
 import logging
 from types import FunctionType
 from typing import Any, Callable, Tuple
@@ -133,22 +134,22 @@ def _generic_plugin_desc(*args: Any, **kwargs: Any) -> Tuple[trtp.TensorDesc]:
         shape_calc_fns = [None] * output.ndim
 
         for i in range(output.ndim):
-            input_node_expr = [
-                syms_arg[j].node.expr
-                for syms_arg in syms_args
-                for j in range(len(syms_arg))
-            ]
+            input_node_expr = input_node_expr = list(
+                itertools.chain.from_iterable(
+                    [sym.node.expr for sym in syms_arg] for syms_arg in syms_args
+                )
+            )
+
             shape_calc_fns[i] = lambdify(
                 tuple(input_node_expr), output.shape[i].node.expr, "math"
             )
 
         out_desc = tensor_args[0].like()
         for i in range(out_desc.ndim):
-            input_shape_expr = [
-                arg.shape_expr[j]
-                for arg in tensor_args
-                for j in range(len(arg.shape_expr))
-            ]
+            input_shape_expr = list(
+                itertools.chain.from_iterable(arg.shape_expr for arg in tensor_args)
+            )
+
             if output.shape[i].node.expr is None:
                 raise ValueError(f"output.shape[{i}].node.expr cannot be None")
             out_desc.shape_expr[i] = shape_calc_fns[i](*input_shape_expr)  # type: ignore[misc]