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

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -0,0 +1,185 @@
+# 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.
+# ==============================================================================
+"""Utilities for building MLIR lowerings."""
+
+import numbers
+from typing import Any
+from typing import Optional
+
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import numpy as np
+
+
+def splat(val, ty, shape=tuple(), *, loc: Optional[Any] = None):
+  if isinstance(ty, ir.IntegerType):
+    if ty.width == 1:
+      attr = ir.BoolAttr.get(bool(val))
+    else:
+      attr = ir.IntegerAttr.get(ty, int(val))
+  elif isinstance(ty, ir.FloatType):
+    attr = ir.FloatAttr.get(ty, val)
+  else:
+    raise ValueError("Unsupported type: %s" % str(ty))
+
+  return stablehlo.constant(
+      ir.DenseElementsAttr.get_splat(
+          ir.RankedTensorType.get(shape, ty),
+          attr,
+      ),
+      loc=loc,
+  )
+
+
+def get_common_broadcast_shape(
+    shape_1: list[int], shape_2: list[int]
+) -> Optional[list[int]]:
+  if not shape_1 and not shape_2:
+    return None
+
+  shape_1 = shape_1 if shape_1 else [1]
+  shape_2 = shape_2 if shape_2 else [1]
+
+  length_diff = abs(len(shape_1) - len(shape_2))
+  if len(shape_1) < len(shape_2):
+    shape_1 = [1] * length_diff + shape_1
+  elif len(shape_1) > len(shape_2):
+    shape_2 = [1] * length_diff + shape_2
+
+  common_broadcast_shape = []
+  for idx in reversed(range(len(shape_1))):
+    dim_size1 = shape_1[idx]
+    dim_size2 = shape_2[idx]
+
+    if dim_size1 == dim_size2:
+      common_broadcast_shape.insert(0, dim_size1)
+    elif dim_size1 == 1 or dim_size2 == 1:
+      common_broadcast_shape.insert(0, max(dim_size1, dim_size2))
+    else:
+      return None
+
+  return common_broadcast_shape
+
+
+def get_broadcast_dimensions(
+    shape_from: list[int], shape_to: list[int]
+) -> list[int]:
+  assert get_common_broadcast_shape(shape_from, shape_to) == shape_to
+
+  ret = []
+  for val in range(len(shape_to) - len(shape_from), len(shape_to)):
+    ret.append(val)
+
+  return ir.DenseI64ArrayAttr.get(np.asarray(ret, np.int64))
+
+
+def broadcast_args_if_needed(
+    val_1: ir.Value, val_2: ir.Value
+) -> tuple[Optional[ir.Value], Optional[ir.Value]]:
+  broadcast_shape = get_common_broadcast_shape(
+      val_1.type.shape, val_2.type.shape
+  )
+  if broadcast_shape is None:
+    return None, None
+
+  new_val_1, new_val_2 = val_1, val_2
+
+  if val_1.type.shape != broadcast_shape:
+    new_val_1 = stablehlo.broadcast_in_dim(
+        result=ir.RankedTensorType.get(
+            broadcast_shape, val_1.type.element_type
+        ),
+        operand=val_1,
+        broadcast_dimensions=get_broadcast_dimensions(
+            val_1.type.shape, broadcast_shape
+        ),
+    )
+  if val_2.type.shape != broadcast_shape:
+    new_val_2 = stablehlo.broadcast_in_dim(
+        result=ir.RankedTensorType.get(
+            broadcast_shape, val_2.type.element_type
+        ),
+        operand=val_2,
+        broadcast_dimensions=get_broadcast_dimensions(
+            val_2.type.shape, broadcast_shape
+        ),
+    )
+  return new_val_1, new_val_2
+
+
+def upcast_to_same_type(*vals: ir.Value):
+  if not vals:
+    return None
+  if len(vals) == 1:
+    return vals[0]
+
+  def get_priority(ty: ir.Type):
+    priorities = [
+        ir.IntegerType.get_signless(1),
+        ir.IntegerType.get_signless(16),
+        ir.IntegerType.get_signless(32),
+        ir.IntegerType.get_signless(64),
+        ir.F16Type,
+        ir.F32Type,
+        ir.F64Type,
+    ]
+    for i, tycls in enumerate(priorities):
+      if tycls.isinstance(ty):
+        return i
+    raise ValueError("Unsupported type: %s" % str(ty))
+
+  cast_tycls = type(max([v.type.element_type for v in vals], key=get_priority))
+  new_vals = []
+  for val in vals:
+    if not cast_tycls.isinstance(val.type.element_type):
+      val = stablehlo.convert(
+          ir.RankedTensorType.get(val.type.shape, cast_tycls.get()), val
+      )
+    new_vals.append(val)
+  return tuple(new_vals)
+
+
+def minmax(ty: ir.Type) -> tuple[numbers.Number, numbers.Number]:
+  if isinstance(ty, ir.IntegerType):
+    if ty.is_unsigned:
+      return (0, 1 << ty.width)
+    else:
+      return (-(1 << (ty.width - 1)), (1 << (ty.width - 1)) - 1)
+  elif isinstance(ty, ir.F16Type):
+    return (np.finfo(np.float16).min, np.finfo(np.float16).max)
+  elif isinstance(ty, ir.F32Type):
+    return (np.finfo(np.float32).min, np.finfo(np.float32).max)
+  elif isinstance(ty, ir.F64Type):
+    return (np.finfo(np.float64).min, np.finfo(np.float64).max)
+  else:
+    raise ValueError("Unsupported type: %s" % ty)
+
+
+def convert_int_to_float(t: ir.Value) -> ir.Value:
+  """Converts an input with type ir.IntegerType to an ir.FloatType of equivalent width."""
+  elty = t.type.element_type
+  if not isinstance(elty, ir.IntegerType):
+    raise ValueError(
+        "Expected input with integer type, received %s" % type(elty)
+    )
+
+  if elty.width == 32:
+    return stablehlo.convert(
+        ir.RankedTensorType.get(t.type.shape, ir.F32Type.get()), t
+    )
+  elif elty.width == 64:
+    return stablehlo.convert(
+        ir.RankedTensorType.get(t.type.shape, ir.F64Type.get()), t
+    )

ai_edge_torch/odml_torch/passes/__init__.py

@@ -0,0 +1,38 @@
+# 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 jax._src.lib.mlir import ir
+from jax._src.lib.mlir import passmanager
+
+
+def run_pass(pipeline, module: ir.Module):
+  pm = passmanager.PassManager.parse(pipeline)
+  pm.run(module.operation)
+  return module
+
+
+def canonicalize(module: ir.Module):
+  return run_pass("builtin.module(canonicalize)", module)
+
+
+def cse(module: ir.Module):
+  return run_pass("builtin.module(cse)", module)
+
+
+def inline(module: ir.Module):
+  return run_pass("builtin.module(inline)", module)
+
+
+def strip_debuginfo(module: ir.Module):
+  return run_pass("builtin.module(strip-debuginfo)", module)

ai_edge_torch/odml_torch/tf_integration.py

@@ -0,0 +1,194 @@
+# 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.
+# ==============================================================================
+"""APIs to convert lowered MLIR from PyTorch to TensorFlow and TFLite artifacts."""
+
+import re
+import tempfile
+
+import tensorflow as tf
+import torch
+
+from tensorflow.compiler.tf2xla.python import xla as tfxla
+
+from . import export
+from . import export_utils
+
+
+def torch_dtype_to_tf(dtype):
+  return {
+      torch.double: tf.float64,
+      torch.float32: tf.float32,
+      torch.half: tf.float16,
+      torch.long: tf.int64,
+      torch.int32: tf.int32,
+      torch.int16: tf.int16,
+      torch.bool: tf.bool,
+  }.get(dtype)
+
+
+def _get_shape_with_dynamic(signature: export.VariableSignature):
+  return [
+      None if export_utils.is_torch_dynamic(s) else s for s in signature.shape
+  ]
+
+
+def _mangle_tf_root_scope_name(name):
+  r"""Build the mangled name for tf.Variable.
+
+  TF has more restricted constrain on the variable names at root scope. Root
+  scope name constrain: [A-Za-z0-9.][A-Za-z0-9_.\\-/]* Non-root scope name
+  constrain: [A-Za-z0-9_.\\-/]*
+  https://github.com/tensorflow/tensorflow/blob/51b601fa6bb7e801c0b6ae73c25580e40a8b5745/tensorflow/python/framework/ops.py#L3301-L3302
+  The state_dict key doesn't have such constrain, the name need to be mangled
+  when a root-scoped TF variable is created.
+
+  FX Graph Node may contain characters other than [A-Za-z0-9_.\\-/], replace
+  offending characters with '_'.
+
+  Args:
+    name: the tensor name to be mangled.
+
+  Returns:
+    Mangled name in str.
+  """
+  if name[0] in "._\\-/":
+    name = "k" + name
+  name = re.sub(r"[^^\w\-/\\]+", "_", name)
+  return name
+
+
+def _build_tf_state_dict(
+    lowered: export.MlirLowered,
+) -> dict[str, tf.Variable]:
+  """Build a dictionary of tf.Variable from the state_dict in lowered."""
+  tf_state_dict = {}
+  for sig in lowered.input_signature:
+    if sig.input_spec.is_parameter:
+      name = sig.input_spec.name
+      tf_state_dict[name] = tf.Variable(
+          lowered.state_dict[name].detach().numpy(),
+          trainable=False,
+          name=_mangle_tf_root_scope_name(name),
+      )
+  return tf_state_dict
+
+
+def _extract_call_args(
+    lowered: export.MlirLowered,
+    args,
+    tf_state_dict: dict[str, tf.Variable],
+):
+  """Extract the flattened inputs to built tf.function."""
+  call_args = []
+  for sig in lowered.input_signature:
+    if sig.input_spec.is_user_input:
+      call_args.append(args[sig.input_spec.i])
+    elif sig.input_spec.is_parameter:
+      name = sig.input_spec.name
+      call_args.append(tf_state_dict[name])
+  return call_args
+
+
+def _wrap_as_tf_func(lowered, tf_state_dict):
+  """Build tf.function from lowered and tf_state_dict."""
+
+  def inner(*args):
+    t_outs = [torch_dtype_to_tf(sig.dtype) for sig in lowered.output_signature]
+    s_outs = [_get_shape_with_dynamic(sig) for sig in lowered.output_signature]
+    call_args = _extract_call_args(lowered, args, tf_state_dict)
+    return tfxla.call_module(
+        tuple(call_args),
+        version=5,
+        Tout=t_outs,  # dtype information
+        Sout=s_outs,  # Shape information
+        function_list=[],
+        module=lowered.module_bytecode,
+    )
+
+  return inner
+
+
+def _make_input_signatures(
+    lowered: export.MlirLowered,
+) -> list[tf.TensorSpec]:
+  """Build the input signatures in tf.TensorSpec for building tf.function."""
+  user_input_signature = sorted(
+      [sig for sig in lowered.input_signature if sig.input_spec.is_user_input],
+      key=lambda sig: sig.input_spec.i,
+  )
+  tf_signatures = []
+
+  for sig in user_input_signature:
+    shape = _get_shape_with_dynamic(sig)
+    tf_signatures.append(
+        tf.TensorSpec(
+            shape=shape,
+            dtype=torch_dtype_to_tf(sig.dtype),
+            name=f"args_{sig.input_spec.i}",
+        )
+    )
+  return tf_signatures
+
+
+def mlir_to_tf_function(lowered: export.MlirLowered):
+  """Convert the MLIR lowered to a executable tf.function."""
+  tf_state_dict = _build_tf_state_dict(lowered)
+  return tf.function(
+      _wrap_as_tf_func(lowered, tf_state_dict),
+      input_signature=_make_input_signatures(lowered),
+  )
+
+
+def mlir_to_flatbuffer(lowered: export.MlirLowered):
+  """Convert the MLIR lowered to a TFLite flatbuffer binary."""
+  tf_state_dict = _build_tf_state_dict(lowered)
+  signature_names = [tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+  tf_signatures = [_make_input_signatures(lowered)]
+  tf_functions = [_wrap_as_tf_func(lowered, tf_state_dict)]
+
+  tf_module = tf.Module()
+  tf_module.f = []
+
+  for tf_sig, func in zip(tf_signatures, tf_functions):
+    tf_module.f.append(
+        tf.function(
+            func,
+            input_signature=tf_sig,
+        )
+    )
+
+    tf_module._variables = list(tf_state_dict.values())
+
+  tf_concrete_funcs = [
+      func.get_concrete_function(*tf_sig)
+      for func, tf_sig in zip(tf_module.f, tf_signatures)
+  ]
+
+  # We need to temporarily save since TFLite's from_concrete_functions does not
+  # allow providing names for each of the concrete functions.
+  with tempfile.TemporaryDirectory() as temp_dir_path:
+    tf.saved_model.save(
+        tf_module,
+        temp_dir_path,
+        signatures={
+            sig_name: tf_concrete_funcs[idx]
+            for idx, sig_name in enumerate(signature_names)
+        },
+    )
+
+    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
+    tflite_model = converter.convert()
+
+  return tflite_model

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20240827"
+__version__ = "0.3.0.dev20240829"

{ai_edge_torch_nightly-0.3.0.dev20240827.dist-info → ai_edge_torch_nightly-0.3.0.dev20240829.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20240827
+Version: 0.3.0.dev20240829
 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.3.0.dev20240827.dist-info → ai_edge_torch_nightly-0.3.0.dev20240829.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=48qP37uHT90YPs4eIUQxCiWVwqGEX3idCUs6mQKvX1U,116
 ai_edge_torch/config.py,sha256=PCd9PVrbUNeVIUDFUCnW4goDWU4bjouK28yMYU6VOi0,877
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=7tox6sdFIlCYPLDYpjFcD8cPTSivURCL_VV6-Dt5Sfc,4910
-ai_edge_torch/version.py,sha256=c4RBJOnG4n0Kx5KGoXFPhG6meBQqqytwtV5tqxvWvJ0,706
+ai_edge_torch/version.py,sha256=OF9oSdUOGcmdEp2HSZmEIeCPlRhL3cpviHc_dExhcX8,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=kcv_QgNgeyDmrqwdzHicGNP68w6zF7GJg7YkMEIXp4Q,3759
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -18,7 +18,7 @@ ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=WKI8V9-V50
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=lxnoH-WGLeiQIF8XjMGodjiZEFTxucl7g05N7MR9OPk,796
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=DIfrWDZ1ufAN_uH-oW3k66jTciY7DlLDAb6UKMN14zE,7528
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_mark.py,sha256=4RyGUwR22bZqkn_TnptenFJodc_Q43f4_SBG7gmTbos,1621
-ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=GcxDxj-5KKWTR5xnRKuhRsb6TDHLCXiPXjGnc_97QXs,12604
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=e_JWgGFOSUI9DUtmod396GNH9uJNd2VBL0DXGjbg-cE,12702
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py,sha256=bsYnudRlXp1PJlu4GF25KSogSkBGQPSaecBrUTONKaw,1031
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=HXTDEP6_Z0I0s58H6I0yHz9qrkOxptIjKhxywfe8F80,10637
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/utils.py,sha256=YLMttMg5PdvXTtQ8lxpKb434UGVvYVALV1-xeuH4UGc,2131
@@ -26,16 +26,16 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=mzfL9cf0qBnpmxM_OlMQFvQsEZV2B_Mia9yEJV4J7rI,7135
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=tvj7fWHHmA9ddtcu-Fp3lJ6emaAQMrtK9wCG0cjgRAo,14413
-ai_edge_torch/_convert/test/test_convert_composites.py,sha256=CBiOqq-m7QT2ggBI1jBl9MkTIT5d0nK1tA0BUga0LGs,7994
-ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=0xIkahEU26Qx9GGn6Dm05ObIqJvsCdh692dREcaHEdE,4725
-ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=jLAmyHw5llT2ff8qA8mem3eVN57e_o5EpBnW72ZtP2I,3026
+ai_edge_torch/_convert/test/test_convert.py,sha256=pUYSXuqFg8CAeJ8JkoYf7S0RDLRPVuZUwVOd0xObM6w,14411
+ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
+ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
+ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
 ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
 ai_edge_torch/debug/culprit.py,sha256=7UYVpVWpiCXbMAyThVtHt_kc_poT7sCTh5UUPvcycgk,14832
 ai_edge_torch/debug/utils.py,sha256=vOAL4t6Lj47uhKapfEsc_WHmvwew3eKO9hSJyzvPXnU,1625
 ai_edge_torch/debug/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/debug/test/test_culprit.py,sha256=GjQv4bpz5EVwgxQt7HmpqTzIo_BpsvRmDVWeOmr29HE,3775
-ai_edge_torch/debug/test/test_search_model.py,sha256=3rUSl7pFBfWjK47YhK5B8J1bXrvNhKKIEuNDNfFShHc,1670
+ai_edge_torch/debug/test/test_culprit.py,sha256=SLX4rC-5Dlna8MWHhGRNe72K71AHTFufDrWLlFQn50c,3773
+ai_edge_torch/debug/test/test_search_model.py,sha256=-RuU0QsjqkfzZF2IbeA55MoeVOawhbgiSEu96PmioPE,1668
 ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -64,7 +64,7 @@ ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=7
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=slieF2-QcDCwd4DRZ7snsZIphT97IXpp4plRRsRSwL8,13983
 ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=7oUIJ6HO0vmlhFdkXpqGm9KTB-eM4Ob9VrHSDlIGFOg,30926
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=CAPsW84A8f00nS6fLFeh_XUjCPsDCA5UxHOUsMrLfSU,3450
-ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=Wc94X_kEnbInTAXFgf-VuCvv1A0HxxWrFZ7Tsq3Z8n0,8662
+ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=x9lEEENGNbpx6VTf_LTVudd9d6bs9tLvFUKTl252zEY,8623
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=xychak9hdLd6ieXBYEwrK2BkF8NRZWZSSCijIsESpBA,3420
 ai_edge_torch/generative/examples/stable_diffusion/util.py,sha256=XIXIB0vCvQKOGyIyiZeiIA5DLeSXjkudywvJS4FK7AM,2431
 ai_edge_torch/generative/examples/stable_diffusion/samplers/__init__.py,sha256=uQWKzCD_49ackNFrt50H04dkDXxfAwUCtMWWQre5SVE,830
@@ -109,10 +109,10 @@ ai_edge_torch/generative/quantize/supported_schemes.py,sha256=FjdycEOvxRgBmQdZVu
 ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/quantize/ai_edge_quantizer_glue/translate_recipe.py,sha256=sSHc_4hUEvi-3KmqbpqWbrRKBjCI1AOctM3dr2EH3vk,5263
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/test/test_experimental_ekv.py,sha256=T5-O2RVLJTH7v9w1_uBfp-Y7o3sdGzYq2Tj2wLRNHyI,4357
-ai_edge_torch/generative/test/test_loader.py,sha256=1ZqAq0HY5uIioumsReOVIsbGBx0WkYcl18PvttdJKrk,3381
-ai_edge_torch/generative/test/test_model_conversion.py,sha256=52ciFy_Qol2Xuym6P6EqdL29oai35LSWGvsUwyEdFTo,8477
-ai_edge_torch/generative/test/test_quantize.py,sha256=3SmJm7Kq98gAneU6IGwwJrJYCVH1qwWR6oUxPfb6qiI,5346
+ai_edge_torch/generative/test/test_experimental_ekv.py,sha256=8qv_eVtJW9GPvBEf2hPQe3tpdJ33XShya6MCX1FqrZM,4355
+ai_edge_torch/generative/test/test_loader.py,sha256=_y5EHGgoNOmCuYonsB81UJScHVsTAQXUVd44czMAw6k,3379
+ai_edge_torch/generative/test/test_model_conversion.py,sha256=-RBTQSERP4szm8s8s_WRmGF3mWZA5E2w2QNtl2MqORw,8475
+ai_edge_torch/generative/test/test_quantize.py,sha256=JEsk9SAkHK0SFm44K_quISc5yBBS6yvtBP1MDyFHdFw,5344
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
 ai_edge_torch/generative/utilities/loader.py,sha256=QFZ2lkeoYQ9MZ1CAFVxBHG4OT192SH74UtJCvbDsdeI,12727
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=pKp3AMSbS3otCvgwJRF5M1l4JRNKk-aCKimXzIMSrds,35679
@@ -122,14 +122,38 @@ ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=cjTprggj_cuktSCm7-A25e7Shop3k
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge7rHq1NhCkBIE,1936
 ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=uiYRfzD1T8deCEAGfdAFusRbI41m14zeTt0Lz5lNT3M,9808
 ai_edge_torch/hlfb/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=e53YNSO2w7Sd9Y717jAr6WKjnXq34Tx_52hXRGtGs3A,4833
-ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py,sha256=7Qbba7GJCBc-J1TUwWIvrpBK0Hwza9nift7sKpW2YVE,8449
+ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=ivq0eVjuf31idfNY0E12F4FxdkSI9hwYXapLJBkIf8Q,4831
+ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py,sha256=j8WpeS-mz3Zr4I7p7NwanQzkQNeH0asZ7lz5y7twgQ4,8447
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
 ai_edge_torch/lowertools/_shim.py,sha256=ilL7x1ebUBj1clg7bagrX4y_nVSHiGrvDrOVfuTeenE,3039
 ai_edge_torch/lowertools/common_utils.py,sha256=Z7p-ivOHtddktpnHrlDm_dSoTxJOdEjFXIGQbzjgwQo,4504
-ai_edge_torch/lowertools/odml_torch_utils.py,sha256=EA2ylE4abCyC1G3lYuNouPGx0lmwtIZe7c42dtX0-3g,7146
-ai_edge_torch/lowertools/test_utils.py,sha256=vsjaX3Ix2U1163jVUNSJgK9io2WNUtJjRvNFE9DrqF4,1932
+ai_edge_torch/lowertools/odml_torch_utils.py,sha256=GKfW1X-QSFffQdVlBuD-bNpP265xcdUlfBY3-9I4f_o,7447
+ai_edge_torch/lowertools/test_utils.py,sha256=bPgc2iXX16KYtMNvmsRdKfrCY6UJmcfitfCOvHoD7Oc,1930
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=-SRm9YNsIGsaVd5Cyp2PP-tdLBJH8EDoMFAa2y89a1w,9043
+ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
+ai_edge_torch/odml_torch/_torch_future.py,sha256=jSYHf1CMTJzMizPMbu2b39hAt0ZTR6gQLq67GMe9KTo,2336
+ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
+ai_edge_torch/odml_torch/export.py,sha256=hIGT-JKYbIa6e_G0AD-k4MSTIAMGdHC1hNHMn9CxsYw,10467
+ai_edge_torch/odml_torch/export_utils.py,sha256=q84U69ZQ82hLXw-xncJ8IW-K71Xux-NWlzZTs7hdZWA,5127
+ai_edge_torch/odml_torch/tf_integration.py,sha256=lTFJPPEijLPFmn6qq2jbpVTQOo0YaOTK36kK6rCiyIE,5956
+ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
+ai_edge_torch/odml_torch/composite/mark_tensor.py,sha256=U--rwl-XkWKgkdXCXDn6yySug8FR66o1YFUAIoSaWW4,3523
+ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=2Y52E_gLeoXpMcPpV-svXsgN3JbEIjnPVjm0xkpTUdQ,3319
+ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=9ag6-WWRG50rPCtIV7OpIokEKu2YRyGlMZZqVPWUH6g,762
+ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=1xCXOs3-9UcsOyLFH0uyQwLu7c06iYFTo0NQ7Ckbl2I,1465
+ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
+ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=Jco5zvejxuyl9xHQxZICAKbkgH7x38qPlwUUpD7S15Q,730
+ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=hXvhKtbH7lGytm6QZOKpTmaLJN3kfENBcSIKQ39ReXA,5478
+ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
+ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=GqYk6oBJw7KWeG4_6gxSu_OvYhjJcC2FpGzWPPEdH6w,933
+ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=wV8AUK8dvjLUy3qjqw_IxpiYVDWUMPNZRfi3XYE_hDs,6972
+ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
+ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=B6BILeu-UlwGB1O6g7111X1TaIFznsfxXrB72ygBsBA,3885
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=I0Y4IK7Zap8m6xfxMw7DfQ9Mg4htKOoypdHVAMHqx9c,10669
+ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
+ai_edge_torch/odml_torch/lowerings/registry.py,sha256=dcnxq8vV9rxSQqXkjSg9it7l6oP_sdfH8kIZdQNkQ_4,2653
+ai_edge_torch/odml_torch/lowerings/utils.py,sha256=NczqpsSd3Fn7yVcPC3qllemiZxxDAZgcW1T5l8-W9fE,5593
+ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
 ai_edge_torch/quantize/__init__.py,sha256=aB5dXot04bqyUhpsDFvxt9CIi15QAC4euvqOndJ0XLU,714
 ai_edge_torch/quantize/pt2e_quantizer.py,sha256=CKIEhs9jCcna64qj1jFH9zEbMbRdyeGV_TmSqEBPjes,15741
 ai_edge_torch/quantize/pt2e_quantizer_utils.py,sha256=eARD1LxLi5m7Z0n_psAkeX_AtUp4fNkE--oECBfivv4,36208
@@ -137,8 +161,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 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=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.3.0.dev20240827.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240827.dist-info/METADATA,sha256=bdUw0fORrhThj1hhI4ZjSJpzM5bZc9UG3JqWPYDYfdY,1878
-ai_edge_torch_nightly-0.3.0.dev20240827.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240827.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240827.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/METADATA,sha256=LrexCNdY177vrp17WaGa53bxHH9vuZXT64O5by4HE6Y,1878
+ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20240829.dist-info/RECORD,,