lint fix

Author: davidpissarra

python/mlc_chat/compiler/compile.py

@@ -83,7 +83,7 @@ def _emit_metadata(metadata):
 
 
 def _attach_variable_bounds(mod, model_config):
-    tir_bound_map = dict()
+    tir_bound_map = {}
     tir_bound_map["seq_len"] = model_config.prefill_chunk_size
     if model_config.context_window_size != -1:
         tir_bound_map["total_seq_len"] = model_config.context_window_size

python/mlc_chat/compiler/model/mistral/mistral_model.py

@@ -371,15 +371,21 @@ def _index(x: te.Tensor):  # x[:-1,:]
             b, s, d = x.shape
             return te.compute((b, 1, d), lambda i, _, k: x[i, s - 1, k], name="index")
 
-        hidden_states = self.model(inputs, total_seq_len, rolling_cache_len, kv_seq_len, attention_mask)
+        hidden_states = self.model(
+            inputs, total_seq_len, rolling_cache_len, kv_seq_len, attention_mask
+        )
         hidden_states = op.tensor_expr_op(_index, name_hint="index", args=[hidden_states])
         logits = self.lm_head(hidden_states)
         if logits.dtype != "float32":
             logits = logits.astype("float32")
         return logits
 
     def prefill(
-        self, inputs: Tensor, total_seq_len: tir.Var, rolling_cache_len: tir.Var, kv_seq_len: tir.Var
+        self,
+        inputs: Tensor,
+        total_seq_len: tir.Var,
+        rolling_cache_len: tir.Var,
+        kv_seq_len: tir.Var,
     ):
         """
         Prefilling the prompt.
@@ -428,7 +434,11 @@ def _sliding_window_attention_mask(
         return self.forward(inputs, total_seq_len, rolling_cache_len, kv_seq_len, attention_mask)
 
     def decode(
-        self, inputs: Tensor, total_seq_len: tir.Var, rolling_cache_len: tir.Var, kv_seq_len: tir.Var
+        self,
+        inputs: Tensor,
+        total_seq_len: tir.Var,
+        rolling_cache_len: tir.Var,
+        kv_seq_len: tir.Var,
     ):
         """Decoding step."""
         batch_size, seq_len = inputs.shape