ai-edge-torch-nightly 0.3.0.dev20240827__py3-none-any.whl → 0.3.0.dev20240829__py3-none-any.whl

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -331,7 +331,12 @@ def _aten__native_batch_norm_legit_no_training(node):
   def batch_norm(input, weight, bias, running_mean, running_var, momentum, eps):
     a = input - running_mean
     b = torch.sqrt(running_var + eps)
-    return a / b * weight + bias, None, None
+    out = a / b
+    if weight is not None:
+      out = out * weight
+    if bias is not None:
+      out = out + bias
+    return out, None, None
 
   node.target = batch_norm
 

ai_edge_torch/_convert/test/test_convert.py

@@ -27,7 +27,7 @@ import tensorflow as tf
 import torch
 import torchvision
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 @dataclasses.dataclass

ai_edge_torch/_convert/test/test_convert_composites.py

@@ -21,7 +21,7 @@ from ai_edge_torch.testing import model_coverage
 import parameterized
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 def _func_to_torch_module(func: Callable[..., torch.Tensor]):

ai_edge_torch/_convert/test/test_convert_multisig.py

@@ -17,9 +17,45 @@
 import ai_edge_torch
 from ai_edge_torch.testing import model_coverage
 import torch
-import torchvision
+from torch import nn
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
+
+
+class FullyConnectedModel(nn.Module):
+  """A simple fully connected model with two fully connected layers."""
+
+  def __init__(self, input_size, hidden_size, output_size):
+    super(FullyConnectedModel, self).__init__()
+    self.fc = nn.Linear(input_size, hidden_size)  # Fully connected layer
+    self.relu = nn.ReLU()  # Activation function
+    self.output = nn.Linear(hidden_size, output_size)
+
+  def forward(self, x):
+    x = self.fc(x)
+    x = self.relu(x)
+    x = self.output(x)
+    return x
+
+
+class FullyConvModel(nn.Module):
+  """A simple fully convolutional model with two convolutions."""
+
+  def __init__(self):
+    super(FullyConvModel, self).__init__()
+    self.conv1 = nn.Conv2d(
+        3, 16, kernel_size=3, padding=1
+    )  # Input channels: 3 (RGB), Output channels: 16
+    self.relu = nn.ReLU(inplace=True)
+    self.conv2 = nn.Conv2d(
+        16, 1, kernel_size=1
+    )  # Output channels: 1 (single channel output)
+
+  def forward(self, x):
+    x = self.conv1(x)
+    x = self.relu(x)
+    x = self.conv2(x)
+    return x
 
 
 class TestConvertMultiSignature(googletest.TestCase):
@@ -29,12 +65,12 @@ class TestConvertMultiSignature(googletest.TestCase):
     super().setUp()
     torch.manual_seed(0)
 
-  def test_convert_mobilenet_v2_with_default(self):
+  def test_convert_with_default(self):
     """Tests conversion of a model with two signatures one of which is the default."""
-    torch_module = torchvision.models.mobilenet_v2().eval()
+    torch_module = FullyConvModel().eval()
 
-    args = (torch.randn(4, 3, 224, 224),)
-    large_args = (torch.randn(4, 3, 336, 336),)
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
 
     signature_name = "large_input"
 
@@ -51,12 +87,12 @@ class TestConvertMultiSignature(googletest.TestCase):
         )
     )
 
-  def test_convert_mobilenet_v2_no_default(self):
+  def test_convert_no_default(self):
     """Tests conversion of a model with two signatures none of which is the default."""
-    torch_module = torchvision.models.mobilenet_v2().eval()
+    torch_module = FullyConvModel().eval()
 
-    args = (torch.randn(4, 3, 224, 224),)
-    large_args = (torch.randn(4, 3, 336, 336),)
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
 
     signature_name_1 = "input"
     signature_name_2 = "large_input"
@@ -84,12 +120,12 @@ class TestConvertMultiSignature(googletest.TestCase):
         )
     )
 
-  def test_convert_mobilenet_v2_signature_helper(self):
+  def test_convert_signature_helper(self):
     """Tests the ai_edge_torch.signature helper function works."""
-    torch_module = torchvision.models.mobilenet_v2().eval()
+    torch_module = FullyConvModel().eval()
 
-    args = (torch.randn(4, 3, 224, 224),)
-    large_args = (torch.randn(4, 3, 336, 336),)
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
 
     signature_name = "large_input"
 
@@ -108,39 +144,43 @@ class TestConvertMultiSignature(googletest.TestCase):
 
   def test_convert_separate_modules(self):
     """Tests conversion of two completely different modules as separate signatures."""
-    mobilentv2 = torchvision.models.mobilenet_v2().eval()
-    resnet18 = torchvision.models.resnet18().eval()
+    fully_conv = FullyConvModel().eval()
+    fully_connected = FullyConnectedModel(10, 5, 10).eval()
 
-    mobilenet_args = (torch.randn(4, 3, 224, 224),)
-    resnet_args = (torch.randn(4, 3, 224, 224),)
+    fully_conv_args = (torch.randn(4, 3, 12, 12),)
+    fully_connected_args = (torch.randn(10),)
 
-    mobilenet_signature_name = "mobilentv2"
-    resnet_signature_name = "resnet18"
+    fully_conv_signature_name = "fully_conv"
+    fully_connected_signature_name = "fully_connected"
 
     edge_model = (
         ai_edge_torch.signature(
-            mobilenet_signature_name, mobilentv2, mobilenet_args
+            fully_conv_signature_name, fully_conv, fully_conv_args
+        )
+        .signature(
+            fully_connected_signature_name,
+            fully_connected,
+            fully_connected_args,
         )
-        .signature(resnet_signature_name, resnet18, resnet_args)
-        .convert(resnet18, resnet_args)
+        .convert(fully_connected, fully_connected_args)
     )
 
-    mobilenet_inference_args = (torch.randn(4, 3, 224, 224),)
-    resnet_inference_args = (torch.randn(4, 3, 224, 224),)
+    fully_conv_inference_args = (torch.randn(4, 3, 12, 12),)
+    fully_connected_inference_args = (torch.randn(10),)
     self.assertTrue(
         model_coverage.compare_tflite_torch(
             edge_model,
-            mobilentv2,
-            mobilenet_inference_args,
-            signature_name=mobilenet_signature_name,
+            fully_conv,
+            fully_conv_inference_args,
+            signature_name=fully_conv_signature_name,
         )
     )
     self.assertTrue(
         model_coverage.compare_tflite_torch(
             edge_model,
-            resnet18,
-            resnet_inference_args,
-            signature_name=resnet_signature_name,
+            fully_connected,
+            fully_connected_inference_args,
+            signature_name=fully_connected_signature_name,
         )
     )
 

ai_edge_torch/_convert/test/test_to_channel_last_io.py

@@ -17,7 +17,7 @@
 import ai_edge_torch
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 class Identity(torch.nn.Module):

ai_edge_torch/debug/test/test_culprit.py

@@ -21,7 +21,7 @@ import sys
 from ai_edge_torch.debug import find_culprits
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 _test_culprit_lib = torch.library.Library("test_culprit", "DEF")
 

ai_edge_torch/debug/test/test_search_model.py

@@ -17,7 +17,7 @@
 from ai_edge_torch.debug import _search_model
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 class TestSearchModel(googletest.TestCase):

ai_edge_torch/generative/examples/stable_diffusion/pipeline.py

@@ -15,16 +15,16 @@
 
 import argparse
 import os
-from pathlib import Path
-from typing import Dict, Optional
+import pathlib
+from typing import Optional
 
-import ai_edge_torch.generative.examples.stable_diffusion.samplers as samplers
-from ai_edge_torch.generative.examples.stable_diffusion.tokenizer import Tokenizer  # NOQA
-import ai_edge_torch.generative.examples.stable_diffusion.util as util
-from ai_edge_torch.model import TfLiteModel
+import ai_edge_torch
+from ai_edge_torch.generative.examples.stable_diffusion import samplers
+from ai_edge_torch.generative.examples.stable_diffusion import tokenizer
+from ai_edge_torch.generative.examples.stable_diffusion import util
 import numpy as np
 from PIL import Image
-from tqdm import tqdm
+import tqdm
 
 arg_parser = argparse.ArgumentParser()
 arg_parser.add_argument(
@@ -104,12 +104,12 @@ class StableDiffusion:
       diffusion_ckpt: str,
       decoder_ckpt: str
   ):
-    self.tokenizer = Tokenizer(tokenizer_vocab_dir)
-    self.clip = TfLiteModel.load(clip_ckpt)
-    self.decoder = TfLiteModel.load(decoder_ckpt)
-    self.diffusion = TfLiteModel.load(diffusion_ckpt)
+    self.tokenizer = tokenizer.Tokenizer(tokenizer_vocab_dir)
+    self.clip = ai_edge_torch.model.TfLiteModel.load(clip_ckpt)
+    self.decoder = ai_edge_torch.model.TfLiteModel.load(decoder_ckpt)
+    self.diffusion = ai_edge_torch.model.TfLiteModel.load(diffusion_ckpt)
     if encoder_ckpt is not None:
-      self.encoder = TfLiteModel.load(encoder_ckpt)
+      self.encoder = ai_edge_torch.model.TfLiteModel.load(encoder_ckpt)
 
 
 def run_tflite_pipeline(
@@ -128,48 +128,32 @@ def run_tflite_pipeline(
 ):
   """Run stable diffusion pipeline with tflite model.
 
-  model:
-
-    StableDiffsuion model.
-  prompt:
-    The prompt to guide the image generation.
-  output_path:
-    The path to the generated output image.
-  uncond_prompt:
-    The prompt not to guide the image generation.
-  cfg_scale:
-    Guidance scale of classifier-free guidance. Higher guidance scale encourages
-    to generate
-    images that are closely linked to the text `prompt`, usually at the expense
-    of lower
-    image quality.
-  height:
-    The height in pixels of the generated image.
-  width:
-    The width in pixels of the generated image.
-  sampler:
-    A sampler to be used to denoise the encoded image latents. Can be one of
-    `k_lms, `k_euler`,
-    or `k_euler_ancestral`.
-  n_inference_steps:
-    The number of denoising steps. More denoising steps usually lead to a higher
-    quality image at the
-    expense of slower inference. This parameter will be modulated by `strength`.
-  seed:
-    A seed to make generation deterministic.
-  strength:
-    Conceptually, indicates how much to transform the reference `input_image`.
-    Must be between 0 and 1.
-    `input_image` will be used as a starting point, adding more noise to it the
-    larger the `strength`.
-    The number of denoising steps depends on the amount of noise initially
-    added. When `strength` is 1,
-    added noise will be maximum and the denoising process will run for the full
-    number of iterations
-    specified in `n_inference_steps`. A value of 1, therefore, essentially
-    ignores `input_image`.
-  input_image:
-    Image which is served as the starting point for the image generation.
+  Args:
+    model: StableDiffsuion model.
+    prompt: The prompt to guide the image generation.
+    output_path: The path to the generated output image.
+    uncond_prompt: The prompt not to guide the image generation.
+    cfg_scale: Guidance scale of classifier-free guidance. Higher guidance scale
+      encourages to generate images that are closely linked to the text
+      `prompt`, usually at the expense of lower image quality.
+    height: The height in pixels of the generated image.
+    width: The width in pixels of the generated image.
+    sampler: A sampler to be used to denoise the encoded image latents. Can be
+      one of `k_lms, `k_euler`, or `k_euler_ancestral`.
+    n_inference_steps: The number of denoising steps. More denoising steps
+      usually lead to a higher quality image at the expense of slower inference.
+      This parameter will be modulated by `strength`.
+    seed: A seed to make generation deterministic.
+    strength: Conceptually, indicates how much to transform the reference
+      `input_image`. Must be between 0 and 1. `input_image` will be used as a
+      starting point, adding more noise to it the larger the `strength`. The
+      number of denoising steps depends on the amount of noise initially added.
+      When `strength` is 1, added noise will be maximum and the denoising
+      process will run for the full number of iterations specified in
+      `n_inference_steps`. A value of 1, therefore, essentially ignores
+      `input_image`.
+    input_image: Image which is served as the starting point for the image
+      generation.
   """
   if not 0 < strength < 1:
     raise ValueError('strength must be between 0 and 1')
@@ -202,7 +186,8 @@ def run_tflite_pipeline(
   context = np.concatenate([cond_context, uncond_context], axis=0)
   noise_shape = (1, 4, height // 8, width // 8)
 
-  # Initialization starts from input_image if any, otherwise, starts from a random sampling.
+  # Initialization starts from input_image if any, otherwise, starts from a
+  # random sampling.
   if input_image:
     if not hasattr(model, 'encoder'):
       raise AttributeError(
@@ -210,7 +195,6 @@ def run_tflite_pipeline(
           ' input_image.'
       )
     input_image = input_image.resize((width, height))
-    input_image_np = np.array(input_image).astype(np.float32)
     input_image_np = util.rescale(input_image, (0, 255), (-1, 1))
     input_image_np = util.move_channel(input_image_np, to='first')
     encoder_noise = np.random.normal(size=noise_shape).astype(np.float32)
@@ -223,8 +207,8 @@ def run_tflite_pipeline(
     latents *= sampler.initial_scale
 
   # Diffusion process.
-  timesteps = tqdm(sampler.timesteps)
-  for i, timestep in enumerate(timesteps):
+  timesteps = tqdm.tqdm(sampler.timesteps)
+  for _, timestep in enumerate(timesteps):
     time_embedding = util.get_time_embedding(timestep)
 
     input_latents = latents * sampler.get_input_scale()
@@ -242,7 +226,7 @@ def run_tflite_pipeline(
   images = util.rescale(images, (-1, 1), (0, 255), clamp=True)
   images = util.move_channel(images, to='last')
   if not os.path.exists(output_path):
-    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    pathlib.Path(output_path).parent.mkdir(parents=True, exist_ok=True)
   Image.fromarray(images[0].astype(np.uint8)).save(output_path)
 
 

ai_edge_torch/generative/test/test_experimental_ekv.py

@@ -21,7 +21,7 @@ from ai_edge_torch.generative.layers.experimental import ekv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 class TestExternalKVLayers(googletest.TestCase):

ai_edge_torch/generative/test/test_loader.py

@@ -22,7 +22,7 @@ from ai_edge_torch.generative.utilities import loader as loading_utils
 import safetensors.torch
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 class TestLoader(googletest.TestCase):

ai_edge_torch/generative/test/test_model_conversion.py

@@ -24,7 +24,7 @@ from ai_edge_torch.testing import model_coverage
 import numpy as np
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 class TestModelConversion(googletest.TestCase):

ai_edge_torch/generative/test/test_quantize.py

@@ -28,7 +28,7 @@ from ai_edge_torch.testing import model_coverage
 from parameterized import parameterized
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 class TestVerifyRecipes(googletest.TestCase):

ai_edge_torch/hlfb/test/test_mark_pattern.py

@@ -19,7 +19,7 @@ from ai_edge_torch.hlfb import mark_pattern
 from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
 import torch
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 def _export_stablehlo_mlir(model, args=None):

ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py

@@ -22,7 +22,7 @@ from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 import torch.nn.functional as F
 
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 def _export_stablehlo_mlir(model, args):

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -84,13 +84,17 @@ def _wrap_as_tf_func(
     t_outs = [torch_dtype_to_tf(sig.dtype) for sig in bundle.output_signature]
     s_outs = [_get_shape_with_dynamic(sig) for sig in bundle.output_signature]
     call_args = _extract_call_args(bundle, args, tf_state_dict)
+    # HACK: In OSS, we use MLIR pybinding and StableHLO dialect from JAX's
+    # build, which may not have the same StableHLO version as what used in
+    # TFLite converter. Therefore we always serialize MLIR module in VHLO.
+    # TODO(b/362798610) Build MLIR pybinding in ai-edge-torch release.
     call_module_return = tfxla.call_module(
         tuple(call_args),
         version=5,
         Tout=t_outs,  # dtype information
         Sout=s_outs,  # Shape information
         function_list=[],
-        module=bundle.module_bytecode,
+        module=bundle.module_bytecode_vhlo,
     )
     spec = exported_program.call_spec.out_spec
 

ai_edge_torch/lowertools/test_utils.py

@@ -16,7 +16,7 @@
 import re
 from typing import Optional
 from ai_edge_torch import config
-from tensorflow.python.platform import googletest
+from absl.testing import absltest as googletest
 
 
 def _extract_backend_configs(mlir):

ai_edge_torch/odml_torch/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+# ==============================================================================
+from . import composite
+from . import debuginfo
+from . import export
+from . import export_utils
+from . import lowerings
+from . import passes

ai_edge_torch/odml_torch/_torch_future.py

@@ -0,0 +1,61 @@
+# 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.
+# ==============================================================================
+"""Wrappers for latest torch APIs/utilities to maintain backward compatibility with older torch releases."""
+
+import torch
+from torch.fx import _pytree as fx_pytree
+
+
+def graph_module_flat_inputs(ep: torch.export.ExportedProgram, args, kwargs):
+  """Transform args, kwargs of __call__ to args for graph_module.
+
+  self.graph_module takes stuff from state dict as inputs.
+  The invariant is for ep: ExportedProgram is
+  ep(args, kwargs) ==
+    ep.postprocess(ep.graph_module(ep.graph_module_flat_inputs(args, kwargs)))
+  """
+  if hasattr(ep, "_graph_module_flat_inputs"):
+    return ep._graph_module_flat_inputs(args, kwargs)
+
+  if args is None:
+    args = tuple()
+  if kwargs is None:
+    kwargs = {}
+
+  flat_args = args
+  if (in_spec := ep.call_spec.in_spec) is not None:
+    if (
+        in_spec.type == tuple
+        and len(in_spec.children_specs) == 2
+        and in_spec.children_specs[0].type == tuple
+        and in_spec.children_specs[1].type == dict
+    ):
+      # NOTE: this is the case where in_spec is for both args and kwargs
+      flat_args = fx_pytree.tree_flatten_spec((args, kwargs), in_spec)
+    else:
+      flat_args = fx_pytree.tree_flatten_spec(args, in_spec)
+
+  param_buffer_keys = ep.graph_signature.parameters + ep.graph_signature.buffers
+  param_buffer_values = tuple(ep.state_dict[key] for key in param_buffer_keys)
+
+  if hasattr(ep.graph_signature, "lifted_tensor_constants"):
+    ordered_tensor_constants = tuple(
+        ep.tensor_constants[name]
+        for name in ep.graph_signature.lifted_tensor_constants
+    )
+  else:
+    ordered_tensor_constants = tuple()
+
+  return (*param_buffer_values, *flat_args, *ordered_tensor_constants)

ai_edge_torch/odml_torch/_torch_library.py

@@ -0,0 +1,19 @@
+# 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.
+# ==============================================================================
+"""Torch library for registering ODML Torch custom ops."""
+
+import torch
+
+ODML_TORCH_LIB = torch.library.Library("odml_torch", "DEF")

ai_edge_torch/odml_torch/composite/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+from .mark_tensor import mark_tensor_op
+from .stablehlo_composite_builder import StableHLOCompositeBuilder