chore: linting

Author: Chengzhe Xu

examples/dynamo/torch_export_pg.py

@@ -1,18 +1,21 @@
 import torch
 import torch_tensorrt
-from transformers import PaliGemmaProcessor, PaliGemmaForConditionalGeneration
+from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
 from transformers.image_utils import load_image
 
-
 # 1. Model
 DEVICE = torch.device("cuda:0")
 model_id = "google/paligemma2-3b-pt-224"
 url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg"
 image = load_image(url)
 
-model = PaliGemmaForConditionalGeneration.from_pretrained(
-    model_id, torch_dtype=torch.float16
-).eval().to(DEVICE)
+model = (
+    PaliGemmaForConditionalGeneration.from_pretrained(
+        model_id, torch_dtype=torch.float16
+    )
+    .eval()
+    .to(DEVICE)
+)
 processor = PaliGemmaProcessor.from_pretrained(model_id)
 
 prompt = ""
@@ -21,19 +24,21 @@
 
 # 2. PyTorch
 with torch.inference_mode():
-    pyt_generation = model.generate(**model_inputs, max_new_tokens=100, do_sample=False) #, use_cache=False)
-    # 입력 토큰 이후의 새로 생성된 토큰만 취합니다.
+    pyt_generation = model.generate(
+        **model_inputs, max_new_tokens=100, do_sample=False
+    )  # , use_cache=False)
+    # The newly generated tokens after the input tokens.
     pyt_generation = pyt_generation[0][input_len:]
     pyt_decoded = processor.decode(pyt_generation, skip_special_tokens=True)
     print("=============================")
     print("PyTorch generated text:")
     print(pyt_decoded)
     print("=============================")
 
-# (a) Dummy inputs  
+# (a) Dummy inputs
 batch_size = 1
-dummy_input_ids = model_inputs["input_ids"] 
-dummy_attention_mask = model_inputs["attention_mask"] 
+dummy_input_ids = model_inputs["input_ids"]
+dummy_attention_mask = model_inputs["attention_mask"]
 dummy_pixel_values = model_inputs["pixel_values"]
 
 dummy_inputs = {
@@ -42,15 +47,15 @@
     "pixel_values": dummy_pixel_values,
 }
 
-# (b) Dynamic shape 
+# (b) Dynamic shape
 BATCH = torch.export.Dim("batch", min=1, max=2)
 SEQ_LEN = torch.export.Dim("seq_len", min=1, max=1024)
 dynamic_shapes = {
     "input_ids": {0: BATCH, 1: SEQ_LEN},
     "attention_mask": {0: BATCH, 1: SEQ_LEN},
     "pixel_values": {0: BATCH},
 }
-# (c) ExportedProgram  
+# (c) ExportedProgram
 # torch.export.export(
 #     model,
 #     args=(),
@@ -64,50 +69,53 @@
 import torch.utils._pytree as pytree
 import transformers
 
+
 def flatten_hybridcache(hc: transformers.cache_utils.HybridCache):
     """
-    1) HybridCache 내부의 '텐서'들을 리스트로 모은다.
-    2) 텐서가 아닌 값들은 context(dict)에 담는다.
+    1) Collects all tensors inside HybridCache into a list.
+    2) Stores non-tensor values in the context (dictionary).
     """
-    # 1. 텐서로 취급할 것들: is_sliding, key_cache 전체, value_cache 전체
+    # 1. Tensors: is_sliding, entire key_cache, entire value_cache
     flat_tensors = []
-    flat_tensors.append(hc.is_sliding)               # shape: [num_hidden_layers], bool
-    flat_tensors.extend(hc.key_cache)                # List[Tensor]
-    flat_tensors.extend(hc.value_cache)              # List[Tensor]
+    flat_tensors.append(hc.is_sliding)  # shape: [num_hidden_layers], bool
+    flat_tensors.extend(hc.key_cache)  # List[Tensor]
+    flat_tensors.extend(hc.value_cache)  # List[Tensor]
 
-    # 2. 텐서가 아닌 필드는 context로 저장
+    # 2. Store non-tensor fields in the context
     context = {
         "max_cache_len": hc.max_cache_len,
         "max_batch_size": hc.max_batch_size,
         "head_dim": hc.head_dim,
         "dtype": hc.dtype,
         "num_key_value_heads": hc.num_key_value_heads,
-        # unflatten 시에 key_cache / value_cache를 몇 개씩 떼어낼지 알아야 하므로
-        "num_layers": len(hc.key_cache),  # = len(hc.value_cache) = config.num_hidden_layers
+        "num_layers": len(
+            hc.key_cache
+        ),  # = len(hc.value_cache) = config.num_hidden_layers
     }
 
     return flat_tensors, context
 
 
 def unflatten_hybridcache(flat_tensors, context):
     """
-    flatten_hybridcache에서 분리한 (flat_tensors, context)를 받아
-    다시 HybridCache 객체로 복원하는 함수.
+    Restores a HybridCache object from the (flat_tensors, context) produced by flatten_hybridcache.
     """
     num_layers = context["num_layers"]
 
-    # 1. flat_tensors 파싱
-    #    - 첫 번째 요소가 is_sliding
-    #    - 그 다음 num_layers개: key_cache
-    #    - 그 다음 num_layers개: value_cache
+    # 1. Parse flat_tensors
+    #    - First element is is_sliding
+    #    - Next num_layers elements: key_cache
+    #    - Next num_layers elements: value_cache
     is_sliding = flat_tensors[0]
     key_cache = flat_tensors[1 : 1 + num_layers]
-    value_cache = flat_tensors[1 + num_layers : 1 + 2*num_layers]
+    value_cache = flat_tensors[1 + num_layers : 1 + 2 * num_layers]
 
-    # 2. __new__로 빈 HybridCache 객체 생성 (생성자 __init__은 호출 안 함)
-    hc = transformers.cache_utils.HybridCache.__new__(transformers.cache_utils.HybridCache)
+    # 2. Create an empty HybridCache object using __new__ (without calling __init__)
+    hc = transformers.cache_utils.HybridCache.__new__(
+        transformers.cache_utils.HybridCache
+    )
 
-    # 3. 필요한 필드를 직접 셋팅
+    # 3. Manually set required fields
     hc.max_cache_len = context["max_cache_len"]
     hc.max_batch_size = context["max_batch_size"]
     hc.head_dim = context["head_dim"]
@@ -119,14 +127,13 @@ def unflatten_hybridcache(flat_tensors, context):
 
     return hc
 
-# pytree 등록
+
+# Register with pytree
 pytree.register_pytree_node(
-    transformers.cache_utils.HybridCache,
-    flatten_hybridcache,
-    unflatten_hybridcache
+    transformers.cache_utils.HybridCache, flatten_hybridcache, unflatten_hybridcache
 )
 
-# from torch.export._trace import _export  
+# from torch.export._trace import _export
 # exported_program = _export(
 #     model,
 #     args=(),
@@ -138,6 +145,7 @@ def unflatten_hybridcache(flat_tensors, context):
 
 # torch.export._draft_export.draft_export
 import torch.export._draft_export
+
 exported_program = torch.export._draft_export.draft_export(
     model,
     args=(),
@@ -156,20 +164,16 @@ def unflatten_hybridcache(flat_tensors, context):
     device=DEVICE,
     disable_tf32=True,
     use_explicit_typing=True,
-    use_fp32_acc=True,  # FP32 누적을 사용해 정확도를 보존합니다.
+    use_fp32_acc=True,
 )
 
-# ----------------------------
-# 5. TensorRT 모델로 생성 수행
-# ----------------------------
-# (원래의 모델 입력을 GPU로 이동시킨 후 generate() 호출)
+# Execute generation using TensorRT model
 model_inputs = {k: v.to(DEVICE) for k, v in model_inputs.items()}
 with torch.inference_mode():
-    trt_generation = trt_model.generate(**model_inputs, max_new_tokens=100, do_sample=False)
+    trt_generation = trt_model.generate(
+        **model_inputs, max_new_tokens=100, do_sample=False
+    )
     trt_generation = trt_generation[0][input_len:]
     trt_decoded = processor.decode(trt_generation, skip_special_tokens=True)
     print("TensorRT generated text:")
     print(trt_decoded)
-
-
-# pytree._register_pytree_node(transformers.modeling_outputs.MaskedLMOutput, lambda x: ([x.loss, x.logits], None), lambda values, _: transformers.modeling_outputs.MaskedLMOutput(loss=values[0], logits=values[1]))

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -813,15 +813,11 @@ def aten_ops_select(
 def index_put_validator(
     node: Node, settings: Optional[CompilationSettings] = None
 ) -> bool:
-    
     if None in node.args[1]:
-        _LOGGER.debug(
-            "We do not support None index yet."
-        )
+        _LOGGER.debug("We do not support None index yet.")
         return False
     else:
         return True
-    
 
 
 @dynamo_tensorrt_converter(

py/torch_tensorrt/dynamo/conversion/impl/select.py

@@ -2,6 +2,7 @@
 from typing import Optional, Sequence, Union
 
 import numpy as np
+import tensorrt as trt
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
@@ -13,7 +14,6 @@
     get_positive_dim,
     get_trt_tensor,
     to_numpy,
-    create_constant,
 )
 from torch_tensorrt.dynamo.conversion.impl.elementwise import convert_binary_elementwise
 from torch_tensorrt.dynamo.conversion.impl.shape import shape as get_shape
@@ -24,8 +24,6 @@
 )
 from torch_tensorrt.fx.types import TRTTensor
 
-import tensorrt as trt
-
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
@@ -499,12 +497,12 @@ def index_put_converter(
         ctx, target, source_ir, f"{name}_cat_indices", reshaped_indices, dim=1
     )
 
-    source_shape = tuple(input_tensor.shape)  
+    source_shape = tuple(input_tensor.shape)
     k = len(indices)
-    leftover_dims = source_shape[k:] 
+    leftover_dims = source_shape[k:]
 
-    index_shapes_py = [tuple(idx.shape) for idx in reshaped_indices] 
-    N = index_shapes_py[0][0] 
+    index_shapes_py = [tuple(idx.shape) for idx in reshaped_indices]
+    N = index_shapes_py[0][0]
     sub_tensor_shape = (N,) + leftover_dims
 
     broadcasted_values = impl.slice.expand(

py/torch_tensorrt/dynamo/utils.py

@@ -378,16 +378,18 @@ def extract_var_range_info(symbolic_integer: torch.SymInt) -> Dict[str, int]:
         shape_env.var_to_val
     )
     assert var_range, var_val
-    min_val_, max_val_ = var_range.lower, var_range.upper # int(var_range.lower), int(var_range.upper)
+    min_val_, max_val_ = (
+        var_range.lower,
+        var_range.upper,
+    )  # int(var_range.lower), int(var_range.upper)
 
     if isinstance(var_range.lower, torch.utils._sympy.numbers.NegativeIntInfinity):
-        min_val_ = 1  
+        min_val_ = 1
 
     # if isinstance(var_range.upper, torch.utils._sympy.numbers.IntInfinity):
-    #     max_val_ = 2048 
-
+    #     max_val_ = 2048
     min_val = int(min_val_)
-    max_val = int(max_val_)    
+    max_val = int(max_val_)
 
     # Torchdynamo 0/1 specialization outlier
     min_val = 1 if min_val == 2 else min_val

tests/py/dynamo/conversion/test_index_put_aten.py

@@ -136,26 +136,25 @@ class TestIndexPutConverter(DispatchTestCase):
             ),
             param(
                 test_name="3d_indices_float_broadcast_index",
-                source_tensor=torch.zeros([3, 3, 3], dtype = torch.int32),
+                source_tensor=torch.zeros([3, 3, 3], dtype=torch.int32),
                 indices_tensor=(
-                    torch.tensor([0,1], dtype=torch.int32),
-                    torch.tensor([0,1], dtype=torch.int32),
+                    torch.tensor([0, 1], dtype=torch.int32),
+                    torch.tensor([0, 1], dtype=torch.int32),
                 ),
-                value_tensor=torch.tensor([10],  dtype = torch.int32),
+                value_tensor=torch.tensor([10], dtype=torch.int32),
             ),
             param(
                 test_name="3d_indices_broadcast_1dim",
                 source_tensor=torch.zeros([3, 3, 3], dtype=torch.int32),
                 indices_tensor=(torch.tensor([1], dtype=torch.int32),),
                 value_tensor=torch.tensor([7], dtype=torch.int32),
             ),
-                        param(
+            param(
                 test_name="2d_indices_broadcast_1dim",
                 source_tensor=torch.zeros([4, 4], dtype=torch.int32),
                 indices_tensor=(torch.tensor([1, 3], dtype=torch.int32),),
                 value_tensor=torch.tensor([5], dtype=torch.int32),
             ),
-
             # 예시 4) 4D 소스, 2D 인덱스 → 마지막 2차원 전체 브로드캐스트
             param(
                 test_name="4d_indices_broadcast_2dim",
@@ -169,7 +168,7 @@ class TestIndexPutConverter(DispatchTestCase):
             # param(
             #     test_name="4d_indices_none_none_single_idx",
             #     source_tensor=torch.zeros([1, 2, 5, 3], dtype=torch.int32),
-            #     # None이 들어가면 현재 코드에서 문제가 발생할 수 있음
+            #     # indexing with None is WIP.
             #     indices_tensor=(None, None, torch.tensor([2], dtype=torch.int32)),
             #     value_tensor=torch.tensor(
             #         [[[10, 20, 30],