ai-edge-torch-nightly 0.2.0.dev20240730__py3-none-any.whl → 0.2.0.dev20240805__py3-none-any.whl

ai_edge_torch/debug/culprit.py

@@ -23,14 +23,13 @@ import os
 import sys
 from typing import Any, Callable, Generator, List, Optional, Tuple, Union
 
+import ai_edge_torch
+from ai_edge_torch.debug import utils
 from functorch.compile import minifier as fx_minifier
 import torch
 from torch._functorch import aot_autograd
 import torch.utils._pytree as pytree
 
-import ai_edge_torch
-from ai_edge_torch.debug import utils
-
 _torch_float_dtypes = {
     torch.float32,
     torch.float,
@@ -120,21 +119,29 @@ class Culprit(SearchResult):
     # TODO (b/321263453): Support Python code gen with sample arg tensor values.
     random_inputs = True
 
-    graph_module_code = self.graph_module.print_readable(print_output=False).rstrip()
+    graph_module_code = self.graph_module.print_readable(
+        print_output=False
+    ).rstrip()
 
     input_strs = []
     for value in self.inputs:
       if torch.is_tensor(value):
         if not random_inputs:
-          input_strs.append(f"# size={_get_shape_str(value)}, dtype={value.dtype}")
-          input_strs.append(f"torch.load(io.BytesIO({_tensor_to_buffer(value)})),")
+          input_strs.append(
+              f"# size={_get_shape_str(value)}, dtype={value.dtype}"
+          )
+          input_strs.append(
+              f"torch.load(io.BytesIO({_tensor_to_buffer(value)})),"
+          )
         else:
           input_strs.append(_tensor_to_random_tensor_call(value) + ",")
       else:
         input_strs.append(str(value) + ",")
 
     inputs_code = (
-        "_args = (\n" + "\n".join([" " * 4 + code for code in input_strs]) + "\n)"
+        "_args = (\n"
+        + "\n".join([" " * 4 + code for code in input_strs])
+        + "\n)"
     )
 
     code = graph_module_code + "\n\n" + inputs_code
@@ -157,7 +164,9 @@ class Culprit(SearchResult):
         + "from torch import device\n"
         + "import ai_edge_torch\n\n"
         + definitions
-        + f"\n\n_edge_model = ai_edge_torch.convert({_CULPRIT_GRAPH_MODULE_NAME}().eval(), _args)\n"
+        + "\n\n_edge_model ="
+        f" ai_edge_torch.convert({_CULPRIT_GRAPH_MODULE_NAME}().eval(),"
+        " _args)\n"
     )
     if self._runtime_errors:
       code += "_edge_model(*_args)\n"
@@ -212,7 +221,9 @@ def _normalize_getitem_nodes(fx_gm: torch.fx.GraphModule):
   return fx_gm
 
 
-def _erase_unused_inputs(fx_gm: torch.fx.GraphModule, inputs: Tuple[torch.Tensor]):
+def _erase_unused_inputs(
+    fx_gm: torch.fx.GraphModule, inputs: Tuple[torch.Tensor]
+):
   fx_gm = copy.deepcopy(fx_gm)
   inputs = tuple(inputs)
   args = fx_gm.graph.process_inputs(*inputs)
@@ -316,7 +327,9 @@ def _erase_sub_gm_from_gm(
   return fx_gm, fx_inputs
 
 
-def _normalize_minified_fx_gm(fx_gm: torch.fx.GraphModule, inputs: Tuple[torch.Tensor]):
+def _normalize_minified_fx_gm(
+    fx_gm: torch.fx.GraphModule, inputs: Tuple[torch.Tensor]
+):
   fx_gm, inputs = _erase_unused_inputs(fx_gm, inputs)
   fx_gm = _lift_dead_ops_to_outputs(fx_gm)
   fx_gm, _ = aot_autograd.aot_export_module(fx_gm, inputs, trace_joint=False)
@@ -374,7 +387,8 @@ def _search_model(
       ep = torch.export.export(model, export_args)
     except Exception as err:
       raise ValueError(
-          "Your model is not exportable by torch.export.export. Please modify your model to be torch-exportable first."
+          "Your model is not exportable by torch.export.export. Please modify"
+          " your model to be torch-exportable first."
       ) from err
   else:
     ep = model
@@ -392,7 +406,9 @@ def _search_model(
       xla_hlo_debug_value = os.environ["XLA_HLO_DEBUG"]
       del os.environ["XLA_HLO_DEBUG"]
 
-    create_minified_hlo_graph = torch._functorch.fx_minifier.create_minified_hlo_graph
+    create_minified_hlo_graph = (
+        torch._functorch.fx_minifier.create_minified_hlo_graph
+    )
     torch._functorch.fx_minifier.create_minified_hlo_graph = (
         lambda *args, **kwargs: None
     )
@@ -403,7 +419,9 @@ def _search_model(
       if xla_hlo_debug_value is not None:
         os.environ["XLA_HLO_DEBUG"] = xla_hlo_debug_value
 
-      torch._functorch.fx_minifier.create_minified_hlo_graph = create_minified_hlo_graph
+      torch._functorch.fx_minifier.create_minified_hlo_graph = (
+          create_minified_hlo_graph
+      )
 
   found_culprits_num = 0
   while True:
@@ -420,7 +438,9 @@ def _search_model(
               max_granularity=max_granularity,
           )
 
-      min_fx_gm, min_inputs = _normalize_minified_fx_gm(raw_min_fx_gm, raw_min_inputs)
+      min_fx_gm, min_inputs = _normalize_minified_fx_gm(
+          raw_min_fx_gm, raw_min_inputs
+      )
       found_culprits_num += 1
       yield SearchResult(min_fx_gm, min_inputs)
 
@@ -429,7 +449,10 @@ def _search_model(
       )
 
     except RuntimeError as e:
-      if str(e) == "Input graph did not fail the tester" and found_culprits_num > 0:
+      if (
+          str(e) == "Input graph did not fail the tester"
+          and found_culprits_num > 0
+      ):
         break
       raise e
 
@@ -467,5 +490,7 @@ def find_culprits(
       enable_fx_minifier_logging=enable_fx_minifier_logging,
   ):
     yield Culprit(
-        search_result.graph_module, search_result.inputs, _runtime_errors=runtime_errors
+        search_result.graph_module,
+        search_result.inputs,
+        _runtime_errors=runtime_errors,
     )

ai_edge_torch/debug/test/test_culprit.py

@@ -19,16 +19,17 @@ import io
 import sys
 import unittest
 
-import torch
-
 from ai_edge_torch.debug import find_culprits
+import torch
 
 _test_culprit_lib = torch.library.Library("test_culprit", "DEF")
 
 _test_culprit_lib.define("non_lowerable_op(Tensor x) -> Tensor")
 
 
-@torch.library.impl(_test_culprit_lib, "non_lowerable_op", "CompositeExplicitAutograd")
+@torch.library.impl(
+    _test_culprit_lib, "non_lowerable_op", "CompositeExplicitAutograd"
+)
 def non_lowerable_op(x):
   if x.max() > 10.0:
     return x + 1.0

ai_edge_torch/debug/test/test_search_model.py

@@ -16,9 +16,8 @@
 
 import unittest
 
-import torch
-
 from ai_edge_torch.debug import _search_model
+import torch
 
 
 class TestSearchModel(unittest.TestCase):
@@ -43,7 +42,9 @@ class TestSearchModel(unittest.TestCase):
 
     results = list(_search_model(find_subgraph_with_sub, model, args))
     self.assertEqual(len(results), 2)
-    self.assertIn(torch.ops.aten.sub.Tensor, [n.target for n in results[0].graph.nodes])
+    self.assertIn(
+        torch.ops.aten.sub.Tensor, [n.target for n in results[0].graph.nodes]
+    )
 
 
 if __name__ == "__main__":

ai_edge_torch/debug/utils.py

@@ -21,7 +21,9 @@ import torch.fx._pytree as fx_pytree
 from torch.utils import _pytree as pytree
 
 
-def exported_program_to_fx_graph_module_and_inputs(ep: torch.export.ExportedProgram):
+def exported_program_to_fx_graph_module_and_inputs(
+    ep: torch.export.ExportedProgram,
+):
   fx_gm = ep.graph_module
   fx_inputs = pytree.tree_map(
       torch.tensor, ep._graph_module_flat_inputs(*ep.example_inputs)

ai_edge_torch/generative/examples/experimental/gemma/convert_to_tflite.py

@@ -20,12 +20,11 @@
 import os
 from pathlib import Path
 
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.experimental.gemma import gemma
 from ai_edge_torch.generative.layers.experimental import ekv_cache as kv_utils
 from ai_edge_torch.generative.quantize import quant_recipes
+import torch
 
 
 def convert_gemma_to_tflite(
@@ -79,7 +78,9 @@ def convert_gemma_to_tflite(
       )
       .convert(quant_config=quant_config)
   )
-  edge_model.export(f'/tmp/gemma_seq{prefill_seq_len}_ekv{kv_cache_max_len}.tflite')
+  edge_model.export(
+      f'/tmp/gemma_seq{prefill_seq_len}_ekv{kv_cache_max_len}.tflite'
+  )
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/experimental/gemma/gemma.py

@@ -21,16 +21,15 @@ import os
 from pathlib import Path
 from typing import Tuple
 
-import numpy as np
-import torch
-import torch.nn as nn
-
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.builder as builder
 from ai_edge_torch.generative.layers.experimental import ekv_cache as kv_utils
 from ai_edge_torch.generative.layers.experimental.attention import TransformerBlock  # NOQA
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import numpy as np
+import torch
+import torch.nn as nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.up_proj",
@@ -81,7 +80,9 @@ class Gemma(nn.Module):
         device=torch.device("cpu"),
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
+        size=config.kv_cache_max,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
     )
     self.config = config
 
@@ -93,9 +94,10 @@ class Gemma(nn.Module):
       kv_cache: kv_utils.EKVCache,
   ) -> Tuple[torch.Tensor, kv_utils.EKVCache]:
     B, T = tokens.size()
-    assert (
-        self.config.max_seq_len >= T
-    ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
+    assert self.config.max_seq_len >= T, (
+        f"Cannot forward sequence of length {T}, max seq length is only"
+        f" {self.config.max_seq_len}"
+    )
 
     cos, sin = self.rope_cache
     cos = cos.index_select(0, input_pos)

ai_edge_torch/generative/examples/experimental/phi/convert_to_tflite.py

@@ -19,12 +19,11 @@
 import os
 from pathlib import Path
 
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.experimental.phi import phi2
 from ai_edge_torch.generative.layers.experimental import ekv_cache
 from ai_edge_torch.generative.quantize import quant_recipes
+import torch
 
 
 def convert_phi2_to_tflite(
@@ -46,7 +45,9 @@ def convert_phi2_to_tflite(
       quantize (bool, optional): Whether the model should be quanized.
         Defaults to True.
   """
-  pytorch_model = phi2.build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
+  pytorch_model = phi2.build_model(
+      checkpoint_path, kv_cache_max_len=kv_cache_max_len
+  )
   # Tensors used to trace the model graph during conversion.
   prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
   prefill_input_pos = torch.arange(0, prefill_seq_len)
@@ -76,7 +77,9 @@ def convert_phi2_to_tflite(
       )
       .convert(quant_config=quant_config)
   )
-  edge_model.export(f'/tmp/phi2_seq{prefill_seq_len}_ekv{kv_cache_max_len}.tflite')
+  edge_model.export(
+      f'/tmp/phi2_seq{prefill_seq_len}_ekv{kv_cache_max_len}.tflite'
+  )
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/experimental/phi/phi2.py

@@ -22,16 +22,15 @@ import os
 from pathlib import Path
 from typing import Tuple
 
-import numpy as np
-import torch
-import torch.nn as nn
-
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.builder as builder
 from ai_edge_torch.generative.layers.experimental import ekv_cache as kv_utils
 from ai_edge_torch.generative.layers.experimental.attention import TransformerBlock  # NOQA
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import numpy as np
+import torch
+import torch.nn as nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.fc1",
@@ -76,7 +75,9 @@ class Phi2(nn.Module):
         device=torch.device("cpu"),
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
+        size=config.kv_cache_max,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
     )
     self.config = config
 
@@ -88,9 +89,10 @@ class Phi2(nn.Module):
       kv_cache: kv_utils.EKVCache,
   ) -> Tuple[torch.Tensor, kv_utils.EKVCache]:
     B, T = tokens.size()
-    assert (
-        self.config.max_seq_len >= T
-    ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
+    assert self.config.max_seq_len >= T, (
+        f"Cannot forward sequence of length {T}, max seq length is only"
+        f" {self.config.max_seq_len}"
+    )
 
     cos, sin = self.rope_cache
     cos = cos.index_select(0, input_pos)

ai_edge_torch/generative/examples/experimental/tiny_llama/convert_to_tflite.py

@@ -20,12 +20,11 @@
 import os
 from pathlib import Path
 
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.experimental.tiny_llama import tiny_llama  # NOQA
 from ai_edge_torch.generative.layers.experimental import ekv_cache as kv_utils
 from ai_edge_torch.generative.quantize import quant_recipes
+import torch
 
 
 def convert_tiny_llama_to_tflite(

ai_edge_torch/generative/examples/experimental/tiny_llama/tiny_llama.py

@@ -22,16 +22,15 @@ import os
 from pathlib import Path
 from typing import Tuple
 
-import numpy as np
-import torch
-import torch.nn as nn
-
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.builder as builder
 from ai_edge_torch.generative.layers.experimental import ekv_cache as kv_utils
 from ai_edge_torch.generative.layers.experimental.attention import TransformerBlock  # NOQA
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import numpy as np
+import torch
+import torch.nn as nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.up_proj",
@@ -78,7 +77,9 @@ class TinyLLamma(nn.Module):
         device=torch.device("cpu"),
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
+        size=config.kv_cache_max,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
     )
     self.config = config
 
@@ -90,9 +91,10 @@ class TinyLLamma(nn.Module):
       kv_cache: kv_utils.EKVCache,
   ) -> Tuple[torch.Tensor, kv_utils.EKVCache]:
     B, T = tokens.size()
-    assert (
-        self.config.max_seq_len >= T
-    ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
+    assert self.config.max_seq_len >= T, (
+        f"Cannot forward sequence of length {T}, max seq length is only"
+        f" {self.config.max_seq_len}"
+    )
 
     cos, sin = self.rope_cache
     cos = cos.index_select(0, input_pos)

ai_edge_torch/generative/examples/gemma/convert_to_tflite.py

@@ -16,11 +16,10 @@
 import os
 from pathlib import Path
 
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.gemma import gemma
 from ai_edge_torch.generative.quantize import quant_recipes
+import torch
 
 
 def convert_gemma_to_tflite(
@@ -58,7 +57,9 @@ def convert_gemma_to_tflite(
       .signature('decode', pytorch_model, (decode_token, decode_input_pos))
       .convert(quant_config=quant_config)
   )
-  edge_model.export(f'/tmp/gemma_seq{prefill_seq_len}_kv{kv_cache_max_len}.tflite')
+  edge_model.export(
+      f'/tmp/gemma_seq{prefill_seq_len}_kv{kv_cache_max_len}.tflite'
+  )
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/gemma/gemma.py

@@ -17,15 +17,14 @@
 import os
 from pathlib import Path
 
-import numpy as np
-import torch
-import torch.nn as nn
-
 from ai_edge_torch.generative.layers.attention import TransformerBlock
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.builder as builder
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import numpy as np
+import torch
+import torch.nn as nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.up_proj",
@@ -76,7 +75,9 @@ class Gemma(nn.Module):
         device=torch.device("cpu"),
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
+        size=config.kv_cache_max,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
     )
     self.config = config
 
@@ -86,9 +87,10 @@ class Gemma(nn.Module):
   @torch.inference_mode
   def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
     B, T = idx.size()
-    assert (
-        self.config.max_seq_len >= T
-    ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
+    assert self.config.max_seq_len >= T, (
+        f"Cannot forward sequence of length {T}, max seq length is only"
+        f" {self.config.max_seq_len}"
+    )
 
     cos, sin = self.rope_cache
     cos = cos.index_select(0, input_pos)
@@ -171,7 +173,9 @@ def define_and_run_2b() -> None:
   input_pos = torch.arange(0, kv_cache_max_len)
   lm_logits = model.forward(tokens, input_pos)
   print("comparing with goldens..")
-  assert torch.allclose(gemma_goldens, lm_logits[0, idx.shape[1] - 1, :], atol=1e-05)
+  assert torch.allclose(
+      gemma_goldens, lm_logits[0, idx.shape[1] - 1, :], atol=1e-05
+  )
 
 
 if __name__ == "__main__":

ai_edge_torch/generative/examples/phi2/convert_to_tflite.py

@@ -16,11 +16,10 @@
 import os
 from pathlib import Path
 
-import torch
-
 import ai_edge_torch
 from ai_edge_torch.generative.examples.phi2 import phi2
 from ai_edge_torch.generative.quantize import quant_recipes
+import torch
 
 
 def convert_phi2_to_tflite(
@@ -41,7 +40,9 @@ def convert_phi2_to_tflite(
       quantize (bool, optional): Whether the model should be quanized.
         Defaults to True.
   """
-  pytorch_model = phi2.build_model(checkpoint_path, kv_cache_max_len=kv_cache_max_len)
+  pytorch_model = phi2.build_model(
+      checkpoint_path, kv_cache_max_len=kv_cache_max_len
+  )
   # Tensors used to trace the model graph during conversion.
   prefill_tokens = torch.full((1, prefill_seq_len), 0, dtype=torch.long)
   prefill_input_pos = torch.arange(0, prefill_seq_len)
@@ -56,7 +57,9 @@ def convert_phi2_to_tflite(
       .signature('decode', pytorch_model, (decode_token, decode_input_pos))
       .convert(quant_config=quant_config)
   )
-  edge_model.export(f'/tmp/phi2_seq{prefill_seq_len}_kv{kv_cache_max_len}.tflite')
+  edge_model.export(
+      f'/tmp/phi2_seq{prefill_seq_len}_kv{kv_cache_max_len}.tflite'
+  )
 
 
 if __name__ == '__main__':

ai_edge_torch/generative/examples/phi2/phi2.py

@@ -18,15 +18,14 @@
 import os
 from pathlib import Path
 
-import numpy as np
-import torch
-import torch.nn as nn
-
 from ai_edge_torch.generative.layers.attention import TransformerBlock
 import ai_edge_torch.generative.layers.attention_utils as attn_utils
 import ai_edge_torch.generative.layers.builder as builder
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import numpy as np
+import torch
+import torch.nn as nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
     ff_up_proj="model.layers.{}.mlp.fc1",
@@ -71,7 +70,9 @@ class Phi2(nn.Module):
         device=torch.device("cpu"),
     )
     self.mask_cache = attn_utils.build_causal_mask_cache(
-        size=config.kv_cache_max, dtype=torch.float32, device=torch.device("cpu")
+        size=config.kv_cache_max,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
     )
     self.config = config
 
@@ -81,9 +82,10 @@ class Phi2(nn.Module):
   @torch.inference_mode
   def forward(self, idx: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor:
     B, T = idx.size()
-    assert (
-        self.config.max_seq_len >= T
-    ), f"Cannot forward sequence of length {T}, max seq length is only {self.config.max_seq_len}"
+    assert self.config.max_seq_len >= T, (
+        f"Cannot forward sequence of length {T}, max seq length is only"
+        f" {self.config.max_seq_len}"
+    )
 
     cos, sin = self.rope_cache
     cos = cos.index_select(0, input_pos)
@@ -160,7 +162,9 @@ def define_and_run() -> None:
   input_pos = torch.arange(0, kv_cache_max_len)
   lm_logits = model.forward(tokens, input_pos)
   print("comparing with goldens..")
-  assert torch.allclose(phi2_goldens, lm_logits[0, idx.shape[1] - 1, :], atol=1e-05)
+  assert torch.allclose(
+      phi2_goldens, lm_logits[0, idx.shape[1] - 1, :], atol=1e-05
+  )
 
 
 if __name__ == "__main__":

ai_edge_torch/generative/examples/stable_diffusion/attention.py

@@ -73,7 +73,9 @@ class SelfAttention(nn.Module):
 
 class CrossAttention(nn.Module):
 
-  def __init__(self, n_heads, d_embed, d_cross, in_proj_bias=True, out_proj_bias=True):
+  def __init__(
+      self, n_heads, d_embed, d_cross, in_proj_bias=True, out_proj_bias=True
+  ):
     super().__init__()
     self.q_proj = nn.Linear(d_embed, d_embed, bias=in_proj_bias)
     self.k_proj = nn.Linear(d_cross, d_embed, bias=in_proj_bias)

ai_edge_torch/generative/examples/stable_diffusion/clip.py

@@ -13,25 +13,34 @@
 # limitations under the License.
 # ==============================================================================
 
-import torch
-from torch import nn
-
 from ai_edge_torch.generative.layers.attention import TransformerBlock
 import ai_edge_torch.generative.layers.attention_utils as attention_utils
 import ai_edge_torch.generative.layers.builder as builder
 import ai_edge_torch.generative.layers.model_config as cfg
 import ai_edge_torch.generative.utilities.loader as loading_utils
+import torch
+from torch import nn
 
 TENSOR_NAMES = loading_utils.ModelLoader.TensorNames(
-    ff_up_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.mlp.fc1",
-    ff_down_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.mlp.fc2",
+    ff_up_proj=(
+        "cond_stage_model.transformer.text_model.encoder.layers.{}.mlp.fc1"
+    ),
+    ff_down_proj=(
+        "cond_stage_model.transformer.text_model.encoder.layers.{}.mlp.fc2"
+    ),
     attn_query_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.q_proj",
     attn_key_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.k_proj",
     attn_value_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.v_proj",
     attn_output_proj="cond_stage_model.transformer.text_model.encoder.layers.{}.self_attn.out_proj",
-    pre_attn_norm="cond_stage_model.transformer.text_model.encoder.layers.{}.layer_norm1",
-    pre_ff_norm="cond_stage_model.transformer.text_model.encoder.layers.{}.layer_norm2",
-    embedding="cond_stage_model.transformer.text_model.embeddings.token_embedding",
+    pre_attn_norm=(
+        "cond_stage_model.transformer.text_model.encoder.layers.{}.layer_norm1"
+    ),
+    pre_ff_norm=(
+        "cond_stage_model.transformer.text_model.encoder.layers.{}.layer_norm2"
+    ),
+    embedding=(
+        "cond_stage_model.transformer.text_model.embeddings.token_embedding"
+    ),
     embedding_position="cond_stage_model.transformer.text_model.embeddings.position_embedding.weight",
     final_norm="cond_stage_model.transformer.text_model.final_layer_norm",
     lm_head=None,
@@ -54,7 +63,9 @@ class CLIP(nn.Module):
     self.transformer_blocks = nn.ModuleList(
         TransformerBlock(config) for _ in range(config.num_layers)
     )
-    self.final_norm = builder.build_norm(config.embedding_dim, config.final_norm_config)
+    self.final_norm = builder.build_norm(
+        config.embedding_dim, config.final_norm_config
+    )
 
     self.mask_cache = attention_utils.build_causal_mask_cache(
         size=config.max_seq_len, dtype=torch.float32