ai-edge-torch-nightly 0.3.0.dev20240828__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

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

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/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

ai_edge_torch/odml_torch/composite/mark_tensor.py

@@ -0,0 +1,120 @@
+# 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 json
+from typing import Sequence, Union
+
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
+
+from .. import _torch_library
+from .. import lowerings
+
+CompositeAttrType = dict[
+    str,
+    Union[
+        int,
+        float,
+        bool,
+        str,
+        Sequence[int],
+        Sequence[float],
+        Sequence[bool],
+    ],
+]
+
+
+def _assert_valid_composite_attr(attr: CompositeAttrType):
+  if attr is None:
+    return
+  if not isinstance(attr, dict):
+    raise ValueError("Composite attr must be a Python dictionary.")
+
+  for k, v in attr.items():
+    if not isinstance(k, str):
+      raise ValueError("Composite attr name must be a Python str.")
+
+    invalid_attr_value_error = ValueError(
+        "Composite attr value must be either Python str, float, int, bool,"
+        " list[int], list[float], list[bool]."
+    )
+    if isinstance(v, (list, tuple)):
+      eltys = {type(el) for el in v}
+      if len(eltys) > 1 or next(iter(eltys)) not in (int, float, bool):
+        raise invalid_attr_value_error
+    elif type(v) not in (str, float, int, bool):
+      raise invalid_attr_value_error
+
+
+@torch._dynamo.assume_constant_result
+def serialize_composite_attr(attr: Union[CompositeAttrType, None]):
+  """Serialize the composite attr into a dynamo-tracable value."""
+  if attr is None:
+    return None
+  _assert_valid_composite_attr(attr)
+  return tuple(attr.items())
+
+
+@torch._dynamo.assume_constant_result
+def deserialize_composite_attr(serialized_attr) -> CompositeAttrType:
+  """Deserialize dynamo-tracable composite attribute into its raw value."""
+  if serialized_attr is None:
+    return None
+  return dict(serialized_attr)
+
+
+_torch_library.ODML_TORCH_LIB.define(
+    "mark_tensor(Tensor x, str name, int pos, str id, bool is_input, Any?"
+    " attr=None) -> Tensor"
+)
+
+mark_tensor_op = torch.ops.odml_torch.mark_tensor.default
+
+
+@torch.library.impl(
+    _torch_library.ODML_TORCH_LIB, "mark_tensor", "CompositeExplicitAutograd"
+)
+def mark_tensor(
+    x: torch.Tensor, name: str, pos: int, id: str, is_input: bool, attr=None
+):
+  return x
+
+
+@torch.library.impl(_torch_library.ODML_TORCH_LIB, "mark_tensor", "Meta")
+def mark_tensor_meta(
+    x: torch.Tensor, name: str, pos: int, id: str, is_input: bool, attr=None
+):
+  return torch.empty_like(x)
+
+
+@lowerings.lower(torch.ops.odml_torch.mark_tensor)
+def mark_tensor_lowering(
+    lctx, x: ir.Value, name: str, pos: int, id: str, is_input: bool, attr=None
+):
+  attr = deserialize_composite_attr(attr)
+  return stablehlo.custom_call(
+      [x.type],
+      inputs=[x],
+      call_target_name="mark_tensor",
+      backend_config=ir.StringAttr.get(
+          json.dumps({
+              "name": name,
+              "pos": pos,
+              "id": id,
+              "is_input": is_input,
+              "attr": attr,
+          })
+      ),
+  )

ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py

@@ -0,0 +1,106 @@
+# 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 uuid
+
+import torch
+
+from . import mark_tensor
+
+
+@torch._dynamo.assume_constant_result
+def _get_uuid() -> str:
+  return uuid.uuid4().hex
+
+
+class StableHLOCompositeBuilder:
+  """Builder class for building a StableHLO composite in the lowering."""
+
+  def __init__(self, name: str, attr: mark_tensor.CompositeAttrType = None):
+    """Helper for building a StableHLO Composite by marking input and output tensors.
+
+    It should be used with the StableHLO converters from `torch_xla.stablehlo`.
+
+    Args:
+      name (str): The name of the built StableHLO Composite op.
+      attr (mark_tensor.CompositeAttrType): Attributes of the StableHLO
+        Composite op.
+    """
+
+    self.attr = attr
+    self.name = name
+    self.id = _get_uuid()
+    self._inputs = []
+    self._outputs = []
+
+  def _mark_tensor(self, *tensors: torch.Tensor, is_input: bool):
+    """Mark the input/output tensors of the StableHLO Composite."""
+    marked_tensors = []
+    serialized_attr = (
+        mark_tensor.serialize_composite_attr(self.attr)
+        if not is_input
+        else None
+    )
+
+    for pos, tensor in enumerate(tensors):
+      if not isinstance(tensor, torch.Tensor):
+        raise ValueError(f"input must be a torch tensor. Got {type(tensor)}.")
+      marked_tensors.append(
+          mark_tensor.mark_tensor_op(
+              tensor,
+              name=self.name,
+              pos=pos,
+              id=self.id,
+              is_input=is_input,
+              attr=serialized_attr,
+          )
+      )
+
+    if len(marked_tensors) == 1:
+      return marked_tensors[0]
+    return tuple(marked_tensors)
+
+  def mark_inputs(self, *tensors: torch.Tensor):
+    """Mark the input tensors of the StableHLO Composite.
+
+    This method must only be called once per builder.
+
+    Args:
+      *tensors (torch.Tensor): Torch tensors to mark.
+
+    Returns:
+      marked_tensors (torch.Tensor or Tuple[torch.Tensor]):
+        Torch tensors marked as composite inputs. The tensor inputs of this
+        method
+        should be replaced by the marked tensors in later usages.
+    """
+
+    return self._mark_tensor(*tensors, is_input=True)
+
+  def mark_outputs(self, *tensors: torch.Tensor):
+    """Mark the output tensors of the StableHLO Composite.
+
+    This method must only be called once per builder.
+
+    Args:
+      *tensors (torch.Tensor): Torch tensors to mark.
+
+    Returns:
+      marked_tensors (torch.Tensor or Tuple[torch.Tensor]):
+        Torch tensors marked as composite outputs. The tensor inputs of this
+        method
+        should be replaced by the marked tensors in later usages.
+    """
+
+    return self._mark_tensor(*tensors, is_input=False)

ai_edge_torch/odml_torch/debuginfo/__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 ._build import build_mlir_debuginfo
+from ._op_polyfill import write_mlir_debuginfo_op

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -0,0 +1,43 @@
+# 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
+
+
+def _class_fullname(cls):
+  module = cls.__module__
+  if module == "builtins":
+    return cls.__qualname__
+  return module + "." + cls.__qualname__
+
+
+def _get_hierarchy(node: torch.fx.Node):
+  nn_module_stack = node.meta.get("nn_module_stack", {})
+  layers = []
+  for name, layer in nn_module_stack.values():
+    iid = ("_" + name.split(".")[-1]) if name else ""
+    layer_str = layer if isinstance(layer, str) else _class_fullname(layer)
+    layers.append(layer_str + iid)
+
+  hierachy_str = "/".join(layers) + ";"
+  return hierachy_str
+
+
+def build_mlir_debuginfo(node: torch.fx.Node):
+  """Build the debuginfo string for the given node's lowerings in MLIR."""
+
+  if not hasattr(node, "meta") or "nn_module_stack" not in node.meta:
+    return None
+
+  return _get_hierarchy(node)

ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py

@@ -0,0 +1,55 @@
+# 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.
+# ==============================================================================
+"""Polyfill op for torch.ops.xla.write_mlir_debuginfo.
+
+In odml-torch, MLIR debuginfo is generated in the lowering framework directly
+without the need of an additional torch op to write. This file register a no-op
+placeholder torch op to replace torch.ops.xla.write_mlir_debuginfo in
+ai-edge-torch.
+"""
+
+from jax._src.lib.mlir import ir
+import torch
+
+from .. import _torch_library
+from .. import lowerings
+
+
+_torch_library.ODML_TORCH_LIB.define(
+    "write_mlir_debuginfo(Tensor x, str data) -> Tensor"
+)
+
+write_mlir_debuginfo_op = torch.ops.odml_torch.write_mlir_debuginfo
+
+
+@torch.library.impl(
+    _torch_library.ODML_TORCH_LIB,
+    "write_mlir_debuginfo",
+    "CompositeExplicitAutograd",
+)
+def write_mlir_debuginfo(x: torch.Tensor, _: str):
+  return x
+
+
+@torch.library.impl(
+    _torch_library.ODML_TORCH_LIB, "write_mlir_debuginfo", "Meta"
+)
+def write_mlir_debuginfo_meta(x: torch.Tensor, _: str):
+  return torch.empty_like(x)
+
+
+@lowerings.lower(torch.ops.odml_torch.write_mlir_debuginfo)
+def write_mlir_debuginfo_lowering(lctx, x: ir.Value, _: str):
+  return x