ai-edge-torch-nightly 0.2.0.dev20240620__py3-none-any.whl → 0.2.0.dev20240621__py3-none-any.whl

ai_edge_torch/convert/conversion.py

@@ -30,6 +30,7 @@ from ai_edge_torch.convert.fx_passes import CanonicalizePass
 from ai_edge_torch.convert.fx_passes import InjectMlirDebuginfoPass
 from ai_edge_torch.convert.fx_passes import OptimizeLayoutTransposesPass
 from ai_edge_torch.convert.fx_passes import run_passes
+from ai_edge_torch.generative.fx_passes import run_generative_passes
 from ai_edge_torch.quantize import quant_config as qcfg
 
 os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
@@ -38,6 +39,7 @@ os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
 def _run_convert_passes(
     exported_program: ExportedProgram,
 ) -> ExportedProgram:
+  exported_program = run_generative_passes(exported_program)
   return run_passes(
       exported_program,
       [

ai_edge_torch/generative/fx_passes/__init__.py

@@ -0,0 +1,31 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import torch
+
+from ai_edge_torch.convert.fx_passes import CanonicalizePass
+from ai_edge_torch.convert.fx_passes import run_passes
+from ai_edge_torch.generative.fx_passes.remove_sdpa_zero_mask_pass import RemoveSDPACompositeZeroMaskPass  # NOQA
+
+
+def run_generative_passes(
+    exported_program: torch.export.ExportedProgram,
+) -> torch.export.ExportedProgram:
+  return run_passes(
+      exported_program,
+      [
+          RemoveSDPACompositeZeroMaskPass(),
+          CanonicalizePass(),
+      ],
+  )

ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py

@@ -0,0 +1,47 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import torch
+
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
+
+
+class RemoveSDPACompositeZeroMaskPass(ExportedProgramPassBase):
+
+  def is_zero_tensor_node(self, node: torch.fx.Node):
+    return node.target == torch.ops.aten.zeros.default
+
+  def call(self, exported_program: torch.export.ExportedProgram):
+    graph = exported_program.graph_module.graph
+    for node in graph.nodes:
+      if not (
+          node.op == "call_function"
+          and node.target == torch.ops.xla.mark_tensor.default
+      ):
+        continue
+
+      source, name, io_position, id, is_input = node.args[:5]
+      # Composite info:
+      # - name: odml.scaled_dot_product_attention
+      # - inputs: q, k, v, mask
+      if name == "odml.scaled_dot_product_attention" and is_input and io_position == 3:
+        if self.is_zero_tensor_node(source):
+          # Remove the mark_tensor call on the mask input by
+          # replacing the target with an identity function.
+          node.target = lambda *args, **kwargs: args[0]
+
+    exported_program.graph_module.graph.lint()
+    exported_program.graph_module.recompile()
+    return ExportedProgramPassResult(exported_program, True)

{ai_edge_torch_nightly-0.2.0.dev20240620.dist-info → ai_edge_torch_nightly-0.2.0.dev20240621.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240620
+Version: 0.2.0.dev20240621
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.2.0.dev20240620.dist-info → ai_edge_torch_nightly-0.2.0.dev20240621.dist-info}/RECORD

@@ -1,7 +1,7 @@
 ai_edge_torch/__init__.py,sha256=FPMmuFU3pyMREtjB_san1fy_0PFtAsgA0VZfOYvDrb4,1008
 ai_edge_torch/model.py,sha256=kmcgELjsYl8YzF8nUF6P7q4i8MWS-pLGpfsy-yTUXmE,4243
 ai_edge_torch/convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/convert/conversion.py,sha256=GN2Js232u_5Y118wg3qIfEoYewxbxLl3TpSnO6osi8c,4029
+ai_edge_torch/convert/conversion.py,sha256=8K8jQuaCjlUWoj7jiimxp_zpN6mYThLOcQ858UDcYnE,4159
 ai_edge_torch/convert/conversion_utils.py,sha256=9BqCL38DErv1vEVGtT3BIJVhdwZjw2EQ-_m5UpvVVYE,11294
 ai_edge_torch/convert/converter.py,sha256=bjj5TV5_g4sGyuSh8ThEDydlNMqhkGSY4SzXK6vwhqI,6927
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=EPs4PSIDLuRH5EBETi6deaOvaaf_Q4xD3_9NVcR7x8o,2810
@@ -65,6 +65,8 @@ ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=E4I5OlC4zyl5cxiiu7uTED-zcwYRu210lP1zuT3xLBE,2566
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=IFRLPG9wz_aLl_zV_6CETCjSM03ukA6bZqqyDLVACuw,5651
+ai_edge_torch/generative/fx_passes/__init__.py,sha256=aXvYiaHDvETIrh0Q9DDZA_ZBiazGk80DT6nt7lLtC1o,1172
+ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=IehLwFNwa0C9fnk1pmNmyfuAwwWbuwdyKy46BSqNVdI,1948
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=AW0Qo3uOIe6p1rJNJ6zR_r4fqL2y-6QJHh0yUd-5Yb0,11966
 ai_edge_torch/generative/layers/attention_utils.py,sha256=hXhuyKblPPxKIRzlAf1YNlwHgpbj-6nReRLhRHELx5k,6350
@@ -110,8 +112,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDd
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.2.0.dev20240620.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240620.dist-info/METADATA,sha256=UKHYbU5TpllRKqT1G0WPn58Qv3IwlNVNlzi2qRmJBsw,1748
-ai_edge_torch_nightly-0.2.0.dev20240620.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240620.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240620.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240621.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240621.dist-info/METADATA,sha256=HVK3rk6sQpfrzg_eAnXkOAUGUDEMkBhDP76usC_GVQQ,1748
+ai_edge_torch_nightly-0.2.0.dev20240621.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240621.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240621.dist-info/RECORD,,