diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/compiler/__init__.py b/integrations/tensorflow/bindings/python/pyiree/tf/compiler/__init__.py
index 774412d..ba966db 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/compiler/__init__.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/compiler/__init__.py
@@ -99,7 +99,7 @@
 
 def tf_saved_model_to_compiler_module(
     saved_model_dir: str,
-    exported_names: Collection[str] = (),
+    exported_names: Sequence[str] = (),
     pass_pipeline: Sequence[str] = TF_IMPORT_PASS_PIPELINE,
     compiler_context: Optional[Context] = None) -> Module:
   """Converts a TensorFlow SavedModel into a MLIR module.
@@ -108,8 +108,7 @@
 
   Args:
     saved_model_dir: Directory of the saved model.
-    exported_names: Optional tuple of strings representing the exported names to
-      keep.
+    exported_names: Optional sequence representing the exported names to keep.
     pass_pipeline: Passes to run on the imported module prior to returning.
       Defaults to TF_IMPORT_PASS_PIPELINE.
     compiler_context: The pyiree.compiler.Context() backing the module.
@@ -131,18 +130,17 @@
 
 def compile_tf_saved_model(
     saved_model_dir: str,
-    exported_names: Collection[str] = (),
-    target_backends: Collection[str] = (),
+    exported_names: Sequence[str] = (),
+    target_backends: Sequence[str] = (),
     pass_pipeline: Sequence[str] = TF_IMPORT_PASS_PIPELINE,
     compiler_context: Optional[Context] = None) -> binding.OpaqueBlob:
   """Compiles a TensorFlow SavedModel to IREE in one shot.
 
   Args:
     saved_model_dir: Directory of the saved model.
-    exported_names: Optional tuple of strings representing the exported names to
-      keep.
-    target_backends: The specific target backends to compile for (defaults to
-      all compiled in targets).
+    exported_names: Optional sequence representing the exported names to keep.
+    target_backends: Optional sequence of specific target backends to compile
+      for (defaults to all compiled in targets).
     pass_pipeline: Passes to run on the imported module prior to returning.
       Defaults to TF_IMPORT_PASS_PIPELINE.
     compiler_context: The pyiree.compiler.Context() backing the module.
@@ -160,7 +158,7 @@
 def tf_signature_def_saved_model_to_compiler_module(
     saved_model_dir: str,
     saved_model_tags: Set[str] = set(),
-    exported_names: Collection[str] = [],
+    exported_names: Sequence[str] = (),
     pass_pipeline: Sequence[str] = TF_IMPORT_PASS_PIPELINE,
     compiler_context: Optional[Context] = None) -> Module:
   """Converts a TensorFlow SignatureDef SavedModel into a MLIR module.
@@ -168,8 +166,7 @@
   Args:
     saved_model_dir: Directory of the saved model.
     saved_model_tags: Optional set of tags to use when loading the model.
-    exported_names: Optional tuple of strings representing the exported names to
-      keep.
+    exported_names: Optional sequence representing the exported names to keep.
     pass_pipeline: Passes to run on the imported module prior to returning.
       Defaults to TF_IMPORT_PASS_PIPELINE.
     compiler_context: The pyiree.compiler.Context() backing the module.
@@ -194,8 +191,8 @@
 def compile_tf_signature_def_saved_model(
     saved_model_dir: str,
     saved_model_tags: Set[str] = set(),
-    exported_names: Collection[str] = (),
-    target_backends: Collection[str] = (),
+    exported_names: Sequence[str] = (),
+    target_backends: Sequence[str] = (),
     pass_pipeline: Sequence[str] = TF_IMPORT_PASS_PIPELINE,
     compiler_context: Optional[Context] = None) -> binding.OpaqueBlob:
   """Compiles a TensorFlow SignatureDef SavedModel to IREE in one shot.
@@ -203,10 +200,9 @@
   Args:
     saved_model_dir: Directory of the saved model.
     saved_model_tags: Optional set of tags to use when loading the model.
-    exported_names: Optional tuple of strings representing the exported names to
-      keep.
-    target_backends: The specific target backends to compile for (defaults to
-      all compiled in targets).
+    exported_names: Optional sequence representing the exported names to keep.
+    target_backends: Optional sequence of specific target backends to compile
+      for (defaults to all compiled in targets).
     pass_pipeline: Passes to run on the imported module prior to returning.
       Defaults to TF_IMPORT_PASS_PIPELINE.
     compiler_context: The pyiree.compiler.Context() backing the module.
@@ -222,7 +218,7 @@
 
 def tf_module_to_compiler_module(
     module: tf.Module,
-    exported_names: Collection[str] = (),
+    exported_names: Sequence[str] = (),
     pass_pipeline: Sequence[str] = TF_IMPORT_PASS_PIPELINE,
     compiler_context: Optional[Context] = None,
     saved_model_dir: str = None) -> Module:
@@ -230,8 +226,7 @@
 
   Args:
     module: The tf.Module instance to convert to MLIR
-    exported_names: Optional tuple of strings representing the exported names to
-      keep.
+    exported_names: Optional sequence representing the exported names to keep.
     pass_pipeline: Passes to run on the imported module prior to returning.
       Defaults to TF_IMPORT_PASS_PIPELINE.
     compiler_context: The pyiree.compiler.Context() backing the module.
@@ -259,8 +254,8 @@
 
 
 def compile_tf_module(module: tf.Module,
-                      exported_names: Collection[str] = (),
-                      target_backends: Collection[str] = (),
+                      exported_names: Sequence[str] = (),
+                      target_backends: Sequence[str] = (),
                       pass_pipeline: Sequence[str] = TF_IMPORT_PASS_PIPELINE,
                       compiler_context: Optional[Context] = None,
                       saved_model_dir: str = None):
@@ -268,10 +263,9 @@
 
   Args:
     module: The tf.Module instance to convert to MLIR
-    exported_names: Optional tuple of strings representing the exported names to
-      keep.
-    target_backends: The specific target backends to compile for (defaults to
-      all compiled in targets).
+    exported_names: Optional sequence representing the exported names to keep.
+    target_backends: Optional sequence of specific target backends to compile
+      for (defaults to all compiled in targets).
     pass_pipeline: Passes to run on the imported module prior to returning.
       Defaults to TF_IMPORT_PASS_PIPELINE.
     compiler_context: The pyiree.compiler.Context() backing the module.
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py
index fa3df03..aa9481e 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils.py
@@ -66,22 +66,22 @@
 
 
 def _parse_target_backends() -> Tuple[Sequence[str], Sequence[str]]:
-  """Decodes --target_backends and creates unique names for their artifacts."""
+  """Decodes --target_backends and creates unique ids for them."""
   backend_names = FLAGS.target_backends.split(",")
   backend_to_index = {k: 0 for k in backend_names if backend_names.count(k) > 1}
-  artifact_names = []
+  backend_ids = []
 
   # If there are multiple copies of the same backend_name, index them. e.g.
   # backend_names = ["tf", "iree_vmla", "tf"]
-  # --> artifact_names = ["tf_0", "iree_vmla", "tf_1"]
+  # --> backend_ids = ["tf_0", "iree_vmla", "tf_1"]
   for backend_name in backend_names:
     if backend_name in backend_to_index:
-      artifact_names.append(f"{backend_name}_{backend_to_index[backend_name]}")
+      backend_ids.append(f"{backend_name}_{backend_to_index[backend_name]}")
       backend_to_index[backend_name] += 1
     else:
-      artifact_names.append(backend_name)
+      backend_ids.append(backend_name)
 
-  return backend_names, artifact_names
+  return backend_names, backend_ids
 
 
 def get_target_backends() -> Sequence[tf_utils.BackendInfo]:
@@ -95,10 +95,10 @@
   """
   if FLAGS.target_backends is not None:
     logging.info("Using backends from command line: %s", FLAGS.target_backends)
-    backend_names, names = _parse_target_backends()
+    backend_names, backend_ids = _parse_target_backends()
     backends = [
-        tf_utils.BackendInfo(backend, name)
-        for backend, name in zip(backend_names, names)
+        tf_utils.BackendInfo(backend_name, backend_id)
+        for backend_name, backend_id in zip(backend_names, backend_ids)
     ]
   else:
     # If no backends are specified, use them all.
@@ -261,10 +261,11 @@
       # Extract metadata from module and function.
       self.module_name = module.module_name
       self.compiled_paths = module.compiled_paths
-      self.backend_name = module.backend
+      self.backend_name = module.backend_info.backend_name
+      self.backend_id = module.backend_info.backend_id
+      self.backend_driver = module.backend_info.driver
       self.iree_serializable = module.iree_serializable()
       self.tflite_serializable = module.tflite_serializable()
-      self.backend_driver = module.backend_driver
       self.function_name = function.__name__
       self.function_sourcefile = inspect.getsourcefile(function)
       source, start_line = inspect.getsourcelines(function)
@@ -276,9 +277,10 @@
       self.module_name = _load_dict["module_name"]
       self.compiled_paths = _load_dict["compiled_paths"]
       self.backend_name = _load_dict["backend_name"]
+      self.backend_id = _load_dict["backend_id"]
+      self.backend_driver = _load_dict["backend_driver"]
       self.iree_serializable = _load_dict["iree_serializable"]
       self.tflite_serializable = _load_dict["tflite_serializable"]
-      self.backend_driver = _load_dict["backend_driver"]
       self.function_name = _load_dict["function_name"]
       self.function_sourcefile = _load_dict["function_sourcefile"]
       self.function_line_numbers = _load_dict["function_line_numbers"]
@@ -286,7 +288,7 @@
       self.calls = _load_dict["calls"]
 
   def __str__(self):
-    header = (f"Trace of {self.module_name} compiled to '{self.backend_name}' "
+    header = (f"Trace of {self.module_name} compiled to '{self.backend_id}' "
               f"on function '{self.function_name}':")
     # Give each call a number so it's easier to compare between multiple traces.
     calls = [f"{i + 1}. {str(call)}" for i, call in enumerate(self.calls)]
@@ -327,11 +329,11 @@
 
       if not calls_match:
         logging.error("Comparision between '%s' and '%s' failed on method '%s'",
-                      ref_trace.backend_name, tar_trace.backend_name,
+                      ref_trace.backend_id, tar_trace.backend_id,
                       ref_call.method)
-        logging.error("Reference call '%s':\n%s", ref_trace.backend_name,
+        logging.error("Reference call '%s':\n%s", ref_trace.backend_id,
                       ref_call)
-        logging.error("Target call '%s':\n%s", tar_trace.backend_name, tar_call)
+        logging.error("Target call '%s':\n%s", tar_trace.backend_id, tar_call)
 
       traces_match = traces_match and calls_match
     return traces_match
@@ -434,14 +436,20 @@
       trace_dir: str, path to the directory to serialize the trace to.
     """
 
+    compiled_paths = None
+    if self.compiled_paths is not None:
+      # Convert to a dict to avoid the issues with serializing defaultdicts.
+      compiled_paths = dict(self.compiled_paths)
+
     # Python serialization.
     metadata = {
         "module_name": self.module_name,
-        "compiled_paths": self.compiled_paths,
+        "compiled_paths": compiled_paths,
         "backend_name": self.backend_name,
+        "backend_id": self.backend_id,
+        "backend_driver": self.backend_driver,
         "iree_serializable": self.iree_serializable,
         "tflite_serializable": self.tflite_serializable,
-        "backend_driver": self.backend_driver,
         "function_name": self.function_name,
         "function_sourcefile": self.function_sourcefile,
         "function_line_numbers": self.function_line_numbers,
@@ -493,7 +501,7 @@
 
 
 def _get_trace_dir(artifacts_dir: str, trace: Trace) -> str:
-  trace_dir = os.path.join(artifacts_dir, trace.backend_name, "traces",
+  trace_dir = os.path.join(artifacts_dir, trace.backend_id, "traces",
                            trace.function_name)
   os.makedirs(trace_dir, exist_ok=True)
   return trace_dir
@@ -559,17 +567,17 @@
 def compile_tf_module(
     module_class: Type[tf.Module], exported_names: Sequence[str] = ()
 ) -> Callable[[Any], Any]:
-  """CompiledModuleTestCase decorator that compiles a tf.Module.
-
-  A CompiledModule is created for each backend in --target_backends. They can
-  be accessed individually via self.compiled_modules.backend_name or as a union
-  via self.get_module().
+  """Compiles module_class to each backend that we test.
 
   Args:
     module_class: the tf.Module subclass to compile.
     exported_names: optional iterable of strings representing which of
       module_class's functions to compile. If exported_names is empty all
       functions will be compiled.
+
+  Returns:
+    A 'Modules' namedtuple containing the reference module, target modules and
+    artifacts directory.
   """
 
   # Setup the directory for saving compilation artifacts and traces.
@@ -580,7 +588,7 @@
                                           f"{FLAGS.reference_backend}_ref")
   tar_backend_infos = get_target_backends()
 
-  compile_backend = lambda backend_info: backend_info.compile(
+  compile_backend = lambda backend_info: backend_info.compile_from_class(
       module_class, exported_names, artifacts_dir)
 
   ref_module = compile_backend(ref_backend_info)
@@ -631,7 +639,7 @@
     failed_backend_indices = []
     for i, tar_trace in enumerate(tar_traces):
       logging.info("Comparing the reference backend '%s' with '%s'",
-                   ref_trace.backend_name, tar_trace.backend_name)
+                   ref_trace.backend_id, tar_trace.backend_id)
       traces_match = Trace.compare_traces(ref_trace, tar_trace)
       if not traces_match:
         failed_backend_indices.append(i)
@@ -649,7 +657,7 @@
     if failed_backend_indices:
       # Extract info for logging.
       failed_backends = [
-          tar_traces[i].backend_name for i in failed_backend_indices
+          tar_traces[i].backend_id for i in failed_backend_indices
       ]
       self.fail(
           "Comparision between the reference backend and the following targets "
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils_test.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils_test.py
index 6157c68..e1fb6e3 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils_test.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_test_utils_test.py
@@ -117,8 +117,8 @@
       # Only outputs
       module.get_count()
 
-    module = tf_utils.TfCompiledModule(StatefulCountingModule,
-                                       tf_utils.BackendInfo('tf'))
+    module = tf_utils.TfCompiledModule.create_from_class(
+        StatefulCountingModule, tf_utils.BackendInfo('tf'))
     trace = tf_test_utils.Trace(module, trace_function)
     trace_function(tf_test_utils.TracedModule(module, trace))
 
@@ -140,13 +140,13 @@
       module.increment()
       module.decrement()
 
-    tf_module = tf_utils.TfCompiledModule(StatefulCountingModule,
-                                          tf_utils.BackendInfo('tf'))
+    tf_module = tf_utils.TfCompiledModule.create_from_class(
+        StatefulCountingModule, tf_utils.BackendInfo('tf'))
     tf_trace = tf_test_utils.Trace(tf_module, tf_function)
     tf_function(tf_test_utils.TracedModule(tf_module, tf_trace))
 
-    vmla_module = tf_utils.IreeCompiledModule(StatefulCountingModule,
-                                              tf_utils.BackendInfo('iree_vmla'))
+    vmla_module = tf_utils.IreeCompiledModule.create_from_class(
+        StatefulCountingModule, tf_utils.BackendInfo('iree_vmla'))
     vmla_trace = tf_test_utils.Trace(vmla_module, vmla_function)
     vmla_function(tf_test_utils.TracedModule(vmla_module, vmla_trace))
 
@@ -161,13 +161,13 @@
     def vmla_function(module):
       module.increment_by(np.array([22.], dtype=np.float32))
 
-    tf_module = tf_utils.TfCompiledModule(StatefulCountingModule,
-                                          tf_utils.BackendInfo('tf'))
+    tf_module = tf_utils.TfCompiledModule.create_from_class(
+        StatefulCountingModule, tf_utils.BackendInfo('tf'))
     tf_trace = tf_test_utils.Trace(tf_module, tf_function)
     tf_function(tf_test_utils.TracedModule(tf_module, tf_trace))
 
-    vmla_module = tf_utils.IreeCompiledModule(StatefulCountingModule,
-                                              tf_utils.BackendInfo('iree_vmla'))
+    vmla_module = tf_utils.IreeCompiledModule.create_from_class(
+        StatefulCountingModule, tf_utils.BackendInfo('iree_vmla'))
     vmla_trace = tf_test_utils.Trace(vmla_module, vmla_function)
     vmla_function(tf_test_utils.TracedModule(vmla_module, vmla_trace))
 
@@ -178,12 +178,12 @@
     def trace_function(module):
       module.increment()
       module.increment_by(np.array([81.], dtype=np.float32))
-      module.increment_by_max(
-          np.array([81], dtype=np.float32), np.array([92], dtype=np.float32))
+      module.increment_by_max(np.array([81], dtype=np.float32),
+                              np.array([92], dtype=np.float32))
       module.get_count()
 
-    module = tf_utils.IreeCompiledModule(StatefulCountingModule,
-                                         tf_utils.BackendInfo('iree_vmla'))
+    module = tf_utils.IreeCompiledModule.create_from_class(
+        StatefulCountingModule, tf_utils.BackendInfo('iree_vmla'))
     trace = tf_test_utils.Trace(module, trace_function)
     trace_function(tf_test_utils.TracedModule(module, trace))
 
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
index fc8bab6..b3237f9 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils.py
@@ -16,6 +16,7 @@
 
 # pylint: disable=protected-access
 
+import collections
 import os
 import random
 import re
@@ -91,7 +92,7 @@
 def _setup_mlir_crash_reproducer(
     function: Callable[[Any], Any],
     artifacts_dir: str,
-    backend_name: str,
+    backend_id: str,
 ) -> Callable[[Any], Any]:
   """Wraps `function` so that it a MLIR crash reproducer is saved if it crashes.
 
@@ -100,7 +101,7 @@
   Args:
     function: The callable to decorate.
     artifacts_dir: The directory to write the reproducer to.
-    backend_name: The name of the backend `function` compiles to.
+    backend_id: The unique backend name to use when writting the reproducer.
 
   Returns:
     A function with the same API as the passed function.
@@ -110,7 +111,7 @@
     # Set up a crash reproducer for debugging.
     if artifacts_dir is not None:
       compiler.Context.default_crash_reproducer_path = os.path.join(
-          artifacts_dir, f"reproducer__{backend_name}.mlir")
+          artifacts_dir, f"reproducer__{backend_id}.mlir")
     try:
       results = function(*args, **kwargs)
     except Exception:  # pylint: disable=broad-except
@@ -135,7 +136,7 @@
       MLIR for the module in TF's input dialect.
     iree_input.mlir:
       The MLIR above translated to IREE via compiler.TF_IMPORT_PASS_PIPELINE.
-    backend_name/compiled.vmfb:
+    backend_id/compiled.vmfb:
       A VM FlatBuffer compiled to the target backend from the IREE MLIR above.
 
   Args:
@@ -167,7 +168,7 @@
 
   compiled_path = None
   if artifacts_dir is not None:
-    backend_dir = os.path.join(artifacts_dir, backend_info.name)
+    backend_dir = os.path.join(artifacts_dir, backend_info.backend_id)
     os.makedirs(backend_dir, exist_ok=True)
     compiled_path = os.path.join(backend_dir, "compiled.vmfb")
     logging.info("Saving compiled IREE module to: %s", compiled_path)
@@ -180,7 +181,7 @@
     module: Type[tf.Module],
     backend_info: "BackendInfo",
     exported_names: Sequence[str] = (),
-    artifacts_dir: str = None
+    artifacts_dir: str = None,
 ) -> Tuple[compiler.binding.OpaqueBlob, Union[str, None]]:
   """Compiles a TensorFlow tf.Module and optionally saves compilation artifacts.
 
@@ -193,8 +194,7 @@
   Args:
     module: A tf.Module.
     backend_info: BackendInfo with the details for compiling module to IREE.
-    exported_names: Iterable of dotted function names to consider for
-      compilation.
+    exported_names: Optional sequence representing the exported names to keep.
     artifacts_dir: An optional string pointing to where compilation artifacts
       should be saved. No compilation artifacts will be saved if this is not
       provided.
@@ -212,29 +212,72 @@
                                                 artifacts_dir)
 
   _compile_module = _setup_mlir_crash_reproducer(_compile_module, artifacts_dir,
-                                                 backend_info.name)
+                                                 backend_info.backend_id)
   return _compile_module(module, exported_names, backend_info, artifacts_dir)
 
 
 class CompiledModule(object):
   """Base class for the TF and IREE compiled modules."""
 
-  def __init__(self, module_class: Type[tf.Module], backend_info: "BackendInfo",
-               exported_names: Sequence[str], artifacts_dir: str):
-    """Shared base constructor – not useful on its own."""
-    self._module_class = module_class
-    self._backend_info = backend_info
-    self._exported_names = exported_names
-    self._artifacts_dir = artifacts_dir
+  def __init__(
+      self,
+      module_name: str,
+      backend_info: "BackendInfo",
+      compiled_paths: Dict[str, str],
+  ):
+    """Shared base constructor – not useful on its own.
 
-    # Public attributes:
-    self.backend = self._backend_info.name
-    self.backend_driver = self._backend_info.driver
-    self.module_name = self._module_class.__name__
-    self.compiled_paths = None
+    Args:
+      module_name: A name for this compiled module. In most cases this will be
+        the name of the tf.Module subclass or instance that is compiled.
+      backend_info: BackendInfo with the details about compiling this module.
+      compiled_paths: A dictionary mapping compiled method names to file paths
+        corresponding to their serialized representations.
+    """
+    self.module_name = module_name
+    self.backend_info = backend_info
+    self.compiled_paths = compiled_paths
 
   def reinitialize(self):
-    """Reinitializes to the initial state of the passed module_class."""
+    """Reinitializes all stateful variables."""
+    raise NotImplementedError()
+
+  @classmethod
+  def create_from_class(cls,
+                        module_class: Type[tf.Module],
+                        backend_info: "BackendInfo",
+                        exported_names: Sequence[str] = (),
+                        artifacts_dir: str = None):
+    """Compile a tf.Module subclass to the target backend in backend_info.
+
+    Args:
+      module_class: The tf.Module subclass to compile.
+      backend_info: BackendInfo with the details for compiling module to IREE.
+      exported_names: Optional sequence representing the exported names to keep.
+      artifacts_dir: An optional string pointing to where compilation artifacts
+        should be saved. No compilation artifacts will be saved if this is not
+        provided.
+    """
+    raise NotImplementedError()
+
+  @classmethod
+  def create_from_instance(cls,
+                           module_instance: tf.Module,
+                           backend_info: "BackendInfo",
+                           exported_names: Sequence[str] = (),
+                           artifacts_dir: str = None):
+    """Compile a tf.Module instance to the target backend in backend_info.
+
+    This is only implemented for IreeCompiledModule.
+
+    Args:
+      module_instance: The tf.Module instance to compile.
+      backend_info: BackendInfo with the details for compiling module to IREE.
+      exported_names: Optional sequence representing the exported names to keep.
+      artifacts_dir: An optional string pointing to where compilation artifacts
+        should be saved. No compilation artifacts will be saved if this is not
+        provided.
+    """
     raise NotImplementedError()
 
   def iree_serializable(self):
@@ -244,14 +287,6 @@
     return False
 
 
-def _get_non_inhereted_function_names(cls):
-  """Gets all methods that cls has that its parents don't have."""
-  names = set(dir(cls))
-  for parent in cls.__bases__:
-    names -= set(dir(parent))
-  return list(names)
-
-
 class _FunctionWrapper(object):
 
   def get_serialized_values(self) -> Tuple[Tuple[str], Tuple[str]]:
@@ -277,57 +312,94 @@
 class IreeCompiledModule(CompiledModule):
   """Iree compiled module."""
 
-  def __init__(self,
-               module_class: Type[tf.Module],
-               backend_info: "BackendInfo",
-               exported_names: Sequence[str] = (),
-               artifacts_dir: str = None):
-    """Compile a tf.Module to the target backend in backend_info.
+  def __init__(
+      self,
+      module_name: str,
+      backend_info: "BackendInfo",
+      compiled_paths: Dict[str, str],
+      vm_module: rt.VmModule,
+      config: rt.Config,
+  ):
+    """Base constructor – Use one of the named constructors instead.
 
     Args:
-      module_class: the tf.Module subclass to compile.
-      backend_info: an element of BackendInfo corresponding to the IREE backend
-        to compile to.
-      exported_names: an optional iterable of strings representing which of the
-        module_class's functions to compile. If exported_names is empty all
-        functions will be compiled.
-      artifacts_dir: an optional path to save compilation artifacts to.
+      module_name: A name for this compiled module. In most cases this will be
+        the name of the tf.Module subclass or instance that is compiled.
+      backend_info: BackendInfo with the details about compiling this module.
+      compiled_paths: A dictionary mapping compiled method names to file paths
+        corresponding to their serialized representations.
+      vm_module: A rt.VmModule containing compilation info to wrap.
+      config: A rt.Config containing compilation info to wrap.
     """
-    super().__init__(module_class, backend_info, exported_names, artifacts_dir)
+    super().__init__(module_name, backend_info, compiled_paths)
+    self._vm_module = vm_module
+    self._config = config
+    self.reinitialize()
 
+  @classmethod
+  def create_from_class(cls,
+                        module_class: Type[tf.Module],
+                        backend_info: "BackendInfo",
+                        exported_names: Sequence[str] = (),
+                        artifacts_dir: str = None):
+    """Compile a tf.Module subclass to the target backend in backend_info.
+
+    Args:
+      module_class: The tf.Module subclass to compile.
+      backend_info: BackendInfo with the details for compiling module to IREE.
+      exported_names: Optional sequence representing the exported names to keep.
+      artifacts_dir: An optional string pointing to where compilation artifacts
+        should be saved. No compilation artifacts will be saved if this is not
+        provided.
+    """
     set_random_seed()
-    self._module_blob, compiled_path = _incrementally_compile_tf_module(
-        module=module_class(),
+    module_instance = module_class()
+    return cls.create_from_instance(module_instance, backend_info,
+                                    exported_names, artifacts_dir)
+
+  @classmethod
+  def create_from_instance(cls,
+                           module_instance: tf.Module,
+                           backend_info: "BackendInfo",
+                           exported_names: Sequence[str] = (),
+                           artifacts_dir: str = None):
+    """Compile a tf.Module instance to the target backend in backend_info.
+
+    Args:
+      module_instance: The tf.Module instance to compile.
+      backend_info: BackendInfo with the details for compiling module to IREE.
+      exported_names: Optional sequence representing the exported names to keep.
+      artifacts_dir: An optional string pointing to where compilation artifacts
+        should be saved. No compilation artifacts will be saved if this is not
+        provided.
+    """
+    module_blob, compiled_path = _incrementally_compile_tf_module(
+        module=module_instance,
         backend_info=backend_info,
         exported_names=exported_names,
         artifacts_dir=artifacts_dir)
-    self._module = rt.VmModule.from_flatbuffer(self._module_blob)
-    self._config = rt.Config(driver_name=backend_info.driver)
+    vm_module = rt.VmModule.from_flatbuffer(module_blob)
+    config = rt.Config(driver_name=backend_info.driver)
 
-    self.compiled_paths = None
+    compiled_paths = None
     if compiled_path is not None:
-      if not len(exported_names):
-        # Get all method names on 'module_class' that aren't on 'tf.Module'.
-        # This doesn't address all possbile scenarios.
-        # TODO(meadowlark): Figure out how to get a list of all of the functions
-        # that this module has access to via `pyiree.rt.system_api.BoundModule`.
-        exported_names = _get_non_inhereted_function_names(module_class)
-      self.compiled_paths = dict([
-          (method, compiled_path) for method in exported_names
-      ])
+      # IREE bundles every compiled method into the same compiled module.
+      compiled_paths = collections.defaultdict(lambda: compiled_path)
 
-    self.reinitialize()
+    module_name = type(module_instance).__name__
+
+    return cls(module_name, backend_info, compiled_paths, vm_module, config)
 
   def reinitialize(self):
-    """Reinitializes to the initial state of the passed module_class."""
+    """Reinitializes all stateful variables."""
     # set_random_seed is not needed here because the model_class.__init__ is not
     # called.
-    self._context = rt.SystemContext(modules=[self._module],
+    self._context = rt.SystemContext(modules=[self._vm_module],
                                      config=self._config)
 
   def __getattr__(self, attr: str) -> _IreeFunctionWrapper:
     # Try to resolve it as a function.
-    m = self._context.modules[self._module.name]
+    m = self._context.modules[self._vm_module.name]
     f = m[attr]
     return _IreeFunctionWrapper(self._context, f)
 
@@ -376,29 +448,54 @@
   normalize TensorFlow's output to Numpy.
   """
 
-  def __init__(self,
-               module_class: Type[tf.Module],
-               backend_info: "BackendInfo",
-               exported_names: Sequence[str] = (),
-               artifacts_dir: str = None):
-    """Wrap a tf.Module in a TFCompiledModule facade.
+  def __init__(
+      self,
+      module_name: str,
+      backend_info: "BackendInfo",
+      constructor: Callable[[], tf.Module],
+      exported_names: Sequence[str],
+  ):
+    """Base constructor – Use one of the named constructors instead.
 
     Args:
-      module_class: the tf.Module subclass to 'compile'.
-      backend_info: one of the 'tf*' elements in BackendInfo.
+      module_name: A name for this compiled module. In most cases this will be
+        the name of the tf.Module subclass or instance that is compiled.
+      backend_info: BackendInfo with the details about compiling this module.
+      constructor: A callable (class or function) which returns the tf.Module
+        subclass instance to wrap.
       exported_names: an optional iterable of strings representing which of the
-        module_class's functions should be callable. If exported_names is empty
-        then all functions will be callable.
-      artifacts_dir: an optional path to save compilation artifacts to. Has no
-        effect for this subclass as nothing is compiled.
+        tf.Module subclass instance's functions should be callable. If
+        exported_names is empty then all functions will be callable.
     """
-    super().__init__(module_class, backend_info, exported_names, artifacts_dir)
+    super().__init__(module_name, backend_info, compiled_paths=None)
+    self._constructor = constructor
+    self._exported_names = exported_names
     self.reinitialize()
 
+  @classmethod
+  def create_from_class(cls,
+                        module_class: Type[tf.Module],
+                        backend_info: "BackendInfo",
+                        exported_names: Sequence[str] = (),
+                        artifacts_dir: str = None):
+    """Compile a tf.Module subclass to the target backend in backend_info.
+
+    Args:
+      module_class: The tf.Module subclass to compile.
+      backend_info: BackendInfo with the details for compiling module to IREE.
+      exported_names: Optional sequence representing the exported names to keep.
+      artifacts_dir: An optional string pointing to where compilation artifacts
+        should be saved. No compilation artifacts will be saved if this is not
+        provided.
+    """
+    module_name = module_class.__name__
+    constructor = module_class
+    return cls(module_name, backend_info, constructor, exported_names)
+
   def reinitialize(self):
-    """Reinitializes to the initial state of the passed module_class."""
+    """Reinitializes all stateful variables."""
     set_random_seed()
-    self._tf_module = self._module_class()
+    self._tf_module = self._constructor()
 
   def __getattr__(self, attr: str) -> _TfFunctionWrapper:
     # Try to resolve it as a function.
@@ -412,6 +509,14 @@
     return _TfFunctionWrapper(f)
 
 
+def _get_non_inhereted_function_names(cls):
+  """Gets all methods that cls has that its parents don't have."""
+  names = set(dir(cls))
+  for parent in cls.__bases__:
+    names -= set(dir(parent))
+  return list(names)
+
+
 def _get_concrete_functions(module_class: Type[tf.Module],
                             exported_names: Sequence[str] = ()):
   """Get concrete functions from non-inherited methods or exported_names."""
@@ -425,12 +530,12 @@
   return functions, exported_names
 
 
-def compile_to_tflite(
+def tf_module_to_tflite_interpreters(
     module_class: Type[tf.Module],
     exported_names: Sequence[str] = (),
     artifacts_dir: str = None
 ) -> Tuple[Dict[str, tf.lite.Interpreter], Union[Dict[str, str]], None]:
-  """Compile a dict of TFLite interpreters for the methods on module_class.
+  """Compile a tf.Module to TFLite interpreters for each of its methods.
 
   Args:
     module_class: A tf.Module subclass to compile with TFLite. If module_class
@@ -463,22 +568,14 @@
     if artifacts_dir is not None:
       compiled_paths[name] = tflite_path
 
+  # Convert module_class's methods into TFLite module byte-strings.
   tflite_modules = []
-  names = []
-  if hasattr(module_class, "get_legacy_tflite_saved_model_converter_kwargs"):
-    kwargs = module_class.get_legacy_tflite_saved_model_converter_kwargs()
-    converter = tf.compat.v1.lite.TFLiteConverter.from_saved_model(
-        kwargs["model_path"],
-        input_arrays=kwargs["input_arrays"],
-        output_arrays=kwargs["output_arrays"])
+  functions, names = _get_concrete_functions(module_class, exported_names)
+  for function in functions:
+    converter = tf.lite.TFLiteConverter.from_concrete_functions([function])
     tflite_modules.append(converter.convert())
-    names.append(kwargs["exported_name"])
-  else:
-    functions, names = _get_concrete_functions(module_class, exported_names)
-    for function in functions:
-      converter = tf.lite.TFLiteConverter.from_concrete_functions([function])
-      tflite_modules.append(converter.convert())
 
+  # Load each of the converted methods above into tf.lite.Interpreters.
   for name, tflite_module in zip(names, tflite_modules):
     if artifacts_dir is None:
       with tempfile.TemporaryDirectory() as base_dir:
@@ -537,18 +634,50 @@
 class TfLiteCompiledModule(CompiledModule):
   """Compiles a tf.Module with TFLite and allows it to be called."""
 
-  def __init__(self,
-               module_class: Type[tf.Module],
-               backend_info: "BackendInfo",
-               exported_names: Sequence[str] = (),
-               artifacts_dir: str = None):
-    super().__init__(module_class, backend_info, exported_names, artifacts_dir)
+  def __init__(
+      self,
+      module_name: str,
+      backend_info: "BackendInfo",
+      compiled_paths: Dict[str, str],
+      interpreters: Dict[str, tf.lite.Interpreter],
+  ):
+    """Base constructor – Use one of the named constructors instead.
+
+    Args:
+      module_name: A name for this compiled module. In most cases this will be
+        the name of the tf.Module subclass or instance that is compiled.
+      backend_info: BackendInfo with the details about compiling this module.
+      compiled_paths: A dictionary mapping compiled method names to file paths
+        corresponding to their serialized representations.
+      interpreters: A dict of tf.lite.Interpreters to make callable.
+    """
+    super().__init__(module_name, backend_info, compiled_paths)
+    self._interpreters = interpreters
+
+  @classmethod
+  def create_from_class(cls,
+                        module_class: Type[tf.Module],
+                        backend_info: "BackendInfo",
+                        exported_names: Sequence[str] = (),
+                        artifacts_dir: str = None):
+    """Compile a tf.Module subclass to the target backend in backend_info.
+
+    Args:
+      module_class: The tf.Module subclass to compile.
+      backend_info: BackendInfo with the details for compiling module to IREE.
+      exported_names: Optional sequence representing the exported names to keep.
+      artifacts_dir: An optional string pointing to where compilation artifacts
+        should be saved. No compilation artifacts will be saved if this is not
+        provided.
+    """
     set_random_seed()
-    self._interpreters, self.compiled_paths = compile_to_tflite(
+    interpreters, compiled_paths = tf_module_to_tflite_interpreters(
         module_class, exported_names, artifacts_dir)
+    module_name = module_class.__name__
+    return cls(module_name, backend_info, compiled_paths, interpreters)
 
   def reinitialize(self):
-    """Reinitializes to the initial state of the passed module_class."""
+    """Reinitializes all stateful variables."""
     # This is a noop because TFLite (mostly) doesn't support stateful modules.
     pass
 
@@ -594,36 +723,43 @@
       },
   }
 
-  def __init__(self, backend_name: str, artifact_name: str = None):
+  def __init__(self, backend_name: str, backend_id: str = None):
     """Creates a BackendInfo with the compilation details for backend_name.
 
     Args:
       backend_name: a str specifying which backend to use. Should be one of
         'tf', 'iree_vmla', 'iree_llvmjit', 'iree_vulkan'.
-      artifact_name: an optional str specifying what name to use when saving
-        compiled artifacts.
+      backend_id: an optional str specifying what name to use when saving
+        compiled artifacts. Must satisfy `backend_id.startswith(backend_name)`.
 
     Raises:
       KeyError: if backend_name is not one of ['tf', 'iree_vmla',
       'iree_llvmjit', 'iree_vulkan'].
+      ValueError: if backend_id doesn't start with backend_name.
     """
     if backend_name not in self._name_to_info:
       raise KeyError(
           "Expected backend_name to be one of "
           f"{list(self._name_to_info.keys())} but got '{backend_name}'.")
+    if backend_id is not None and not backend_id.startswith(backend_name):
+      raise ValueError(f"Expected backend_id to start with '{backend_name}' "
+                       f"but got '{backend_id}'.")
+
+    self.backend_name = backend_name
+    self.backend_id = backend_name if backend_id is None else backend_id
+
     info = self._name_to_info[backend_name]
     self._compiled_module_class = info["compiled_module_class"]
     self.driver = info["driver"]
     self.compiler_targets = info["compiler_targets"]
-    self.name = backend_name if artifact_name is None else artifact_name
 
-  def compile(self,
-              module: Type[tf.Module],
-              exported_names: Sequence[str] = (),
-              artifacts_dir: str = None) -> CompiledModule:
+  def compile_from_class(self,
+                         module_class: Type[tf.Module],
+                         exported_names: Sequence[str] = (),
+                         artifacts_dir: str = None) -> CompiledModule:
     """Creates a 'CompiledModule' for this backend."""
-    return self._compiled_module_class(module, self, exported_names,
-                                       artifacts_dir)
+    return self._compiled_module_class.create_from_class(
+        module_class, self, exported_names, artifacts_dir)
 
   @classmethod
   def get_all_backends(cls) -> Sequence["BackendInfo"]:
diff --git a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py
index 1a441d0..a9c259c 100644
--- a/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py
+++ b/integrations/tensorflow/bindings/python/pyiree/tf/support/tf_utils_test.py
@@ -89,7 +89,7 @@
   ])
   def test_unaltered_state(self, backend_name):
     backend_info = tf_utils.BackendInfo(backend_name)
-    module = backend_info.compile(StatefulCountingModule)
+    module = backend_info.compile_from_class(StatefulCountingModule)
 
     # Test that incrementing works properly.
     self.assertEqual([0.], module.get_count())
@@ -126,8 +126,8 @@
     backend_info = tf_utils.BackendInfo(backend_name)
 
     # Test compilation is the same.
-    module_1 = backend_info.compile(RandomInitModule)
-    module_2 = backend_info.compile(RandomInitModule)
+    module_1 = backend_info.compile_from_class(RandomInitModule)
+    module_2 = backend_info.compile_from_class(RandomInitModule)
     self.assertAllEqual(module_1.get(), module_2.get())
 
     # Test reinitialization is the same.
diff --git a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/__init__.py b/integrations/tensorflow/bindings/python/pyiree/xla/compiler/__init__.py
index 78baeb6..df9b317 100644
--- a/integrations/tensorflow/bindings/python/pyiree/xla/compiler/__init__.py
+++ b/integrations/tensorflow/bindings/python/pyiree/xla/compiler/__init__.py
@@ -54,7 +54,7 @@
 def xla_load_module_proto(
     xla_computation,
     compiler_context: Optional[Context] = None,
-    exported_names: Collection[str] = (),
+    exported_names: Sequence[str] = (),
     pass_pipeline: Sequence[str] = XLA_IMPORT_PASS_PIPELINE) -> Module:
   """Loads a XLA saved model from its persistent representation.
 
@@ -63,8 +63,7 @@
   Args:
     xla_computation: XLA Computation generate from XLA Python client
     compiler_context: The pyiree.compiler.Context() backing the module.
-    exported_names: Optional tuple of strings representing the exported names to
-      keep.
+    exported_names: Optional sequence representing the exported names to keep.
     pass_pipeline: Passes to run on the imported module prior to returning.
       Defaults to XLA_IMPORT_PASS_PIPELINE.
 
@@ -85,21 +84,20 @@
 def xla_compile_module_proto(
     xla_computation,
     compiler_context: Optional[Context] = None,
-    exported_names: Collection[str] = (),
+    exported_names: Sequence[str] = (),
     pass_pipeline: Sequence[str] = XLA_IMPORT_PASS_PIPELINE,
-    target_backends: Collection[str] = ()
+    target_backends: Sequence[str] = ()
 ) -> binding.OpaqueBlob:
   """Loads and compiles a XLA saved model in one shot.
 
   Args:
     xla_computation: XLA Computation generate from XLA Python client
     compiler_context: The pyiree.compiler.Context() backing the module.
-    exported_names: Optional tuple of strings representing the exported names to
-      keep.
+    exported_names: Optional sequence representing the exported names to keep.
     pass_pipeline: Passes to run on the imported module prior to returning.
       Defaults to XLA_IMPORT_PASS_PIPELINE.
-    target_backends: The specific target backends to compile for (defaults to
-      all compiled in targets).
+    target_backends: Optional sequence of specific target backends to compile
+      for (defaults to all compiled in targets).
 
   Returns:
     An OpaqueBlob representing the compiled module.
diff --git a/integrations/tensorflow/e2e/BUILD b/integrations/tensorflow/e2e/BUILD
index 2a0b458..dc289d6 100644
--- a/integrations/tensorflow/e2e/BUILD
+++ b/integrations/tensorflow/e2e/BUILD
@@ -95,6 +95,7 @@
     "range_test.py",
     "ring_buffer_test.py",  # TODO(b/148747011)
     "scatter_update_test.py",
+    "sort_test.py",
     "strings_test.py",
 ]
 
@@ -112,6 +113,7 @@
     "range_test.py",
     "ring_buffer_test.py",  # TODO(b/148747011)
     "scatter_update_test.py",
+    "sort_test.py",
     "strings_test.py",
 ]
 
@@ -218,7 +220,6 @@
     name = "mobile_bert_squad_tests",
     backends_to_srcs = {
         "tf": ["mobile_bert_squad_test.py"],
-        "tflite": ["mobile_bert_squad_test.py"],
     },
     reference_backend = "tf",
     tags = [
@@ -234,6 +235,7 @@
 iree_e2e_test_suite(
     name = "mobile_bert_squad_tests_failing",
     backends_to_srcs = {
+        "tflite": ["mobile_bert_squad_test.py"],
         "iree_vmla": ["mobile_bert_squad_test.py"],
         "iree_llvmjit": ["mobile_bert_squad_test.py"],
         "iree_vulkan": ["mobile_bert_squad_test.py"],
diff --git a/iree/compiler/Dialect/HAL/Target/SPIRVCommon/BUILD b/iree/compiler/Dialect/HAL/Target/SPIRVCommon/BUILD
index 9c55ead..ae5d0d2 100644
--- a/iree/compiler/Dialect/HAL/Target/SPIRVCommon/BUILD
+++ b/iree/compiler/Dialect/HAL/Target/SPIRVCommon/BUILD
@@ -22,7 +22,7 @@
 
 iree_cmake_extra_content(
     content = """
-if(NOT ${IREE_TARGET_BACKEND_VULKAN-SPIRV})
+if(NOT ${IREE_TARGET_BACKEND_VULKAN-SPIRV} AND NOT ${IREE_TARGET_BACKEND_METAL-SPIRV})
   return()
 endif()
 """,
diff --git a/iree/compiler/Dialect/HAL/Target/SPIRVCommon/CMakeLists.txt b/iree/compiler/Dialect/HAL/Target/SPIRVCommon/CMakeLists.txt
index 0922d4a..45d0b21 100644
--- a/iree/compiler/Dialect/HAL/Target/SPIRVCommon/CMakeLists.txt
+++ b/iree/compiler/Dialect/HAL/Target/SPIRVCommon/CMakeLists.txt
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-if(NOT ${IREE_TARGET_BACKEND_VULKAN-SPIRV})
+if(NOT ${IREE_TARGET_BACKEND_VULKAN-SPIRV} AND NOT ${IREE_TARGET_BACKEND_METAL-SPIRV})
   return()
 endif()
 
diff --git a/iree/compiler/Dialect/VMLA/Conversion/ConversionTarget.cpp b/iree/compiler/Dialect/VMLA/Conversion/ConversionTarget.cpp
index a80a260..6af672a 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/ConversionTarget.cpp
+++ b/iree/compiler/Dialect/VMLA/Conversion/ConversionTarget.cpp
@@ -42,6 +42,7 @@
   // Pseudo-ops are illegal.
   // If we end up with a lot of these, consider using an "is pseudo" trait.
   addIllegalOp<IREE::VMLA::BatchMatMulPseudoOp>();
+  addIllegalOp<IREE::VMLA::SortPseudoOp>();
 
   // Allow other ops to pass through so long as their type is valid (not a
   // tensor, basically).
diff --git a/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/ConvertHLOToVMLA.cpp b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/ConvertHLOToVMLA.cpp
index ccd1a3f..75b640f 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/ConvertHLOToVMLA.cpp
+++ b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/ConvertHLOToVMLA.cpp
@@ -704,6 +704,29 @@
   TypeConverter &typeConverter;
 };
 
+struct SortOpConversion : public OpConversionPattern<IREE::VMLA::SortPseudoOp> {
+  SortOpConversion(MLIRContext *context, TypeConverter &typeConverter)
+      : OpConversionPattern(context), typeConverter(typeConverter) {}
+
+  LogicalResult matchAndRewrite(
+      IREE::VMLA::SortPseudoOp srcOp, ArrayRef<Value> rawOperands,
+      ConversionPatternRewriter &rewriter) const override {
+    auto inputType =
+        srcOp.getOperand().getType().cast<ShapedType>().getElementType();
+    auto src = rawOperands[0];
+    auto src_shape = VMLAConversionTarget::getTensorShape(
+        srcOp.getLoc(), srcOp.value(), typeConverter, rewriter);
+    auto dst = VMLAConversionTarget::allocateOutputBuffer(
+        srcOp.getLoc(), srcOp.getResult(), typeConverter, rewriter);
+    rewriter.createOrFold<IREE::VMLA::SortOp>(srcOp.getLoc(), src, src_shape,
+                                              dst, TypeAttr::get(inputType));
+    rewriter.replaceOp(srcOp, {dst});
+    return success();
+  }
+
+  TypeConverter &typeConverter;
+};
+
 struct ConvertOpConversion : public OpConversionPattern<mhlo::ConvertOp> {
   ConvertOpConversion(MLIRContext *context, TypeConverter &typeConverter)
       : OpConversionPattern(context), typeConverter(typeConverter) {}
@@ -769,6 +792,9 @@
                                    IREE::VMLA::BatchMatMulOp>>(context,
                                                                typeConverter);
 
+  // vmla.sort.pseudo
+  patterns.insert<SortOpConversion>(context, typeConverter);
+
   // Simple 1:1 conversion patterns using the automated trait-based converter.
   // Used for HLO ops that have equivalent VMLA ops such as most arithmetic ops.
   patterns.insert<VMLAOpConversion<mhlo::AddOp, IREE::VMLA::AddOp>>(
diff --git a/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/test/sort.mlir b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/test/sort.mlir
new file mode 100644
index 0000000..0903793
--- /dev/null
+++ b/iree/compiler/Dialect/VMLA/Conversion/HLOToVMLA/test/sort.mlir
@@ -0,0 +1,37 @@
+// RUN: iree-opt -split-input-file -iree-vmla-pre-conversion-lowering -iree-vmla-conversion -canonicalize %s | IreeFileCheck %s
+
+func @sort1D(%arg0 : tensor<4xf32>) -> tensor<4xf32> attributes { sym_visibility = "private" } {
+  // CHECK-DAG: [[C16:%.+]] = constant 16 : index
+  // CHECK-DAG: [[RS:%.+]] = shapex.const_ranked_shape : !shapex.ranked_shape<[4]>
+  // CHECK-DAG: [[BL:%.+]] = vmla.buffer.alloc byte_length = [[C16]] : !vmla.buffer
+  // CHECK-DAG: vmla.sort %arg0([[RS]] : !shapex.ranked_shape<[4]>), out [[BL]] : f32
+  // CHECK-DAG: [[BUF:%.+]] = vmla.buffer.alloc byte_length = [[C16]] : !vmla.buffer
+  // CHECK-DAG: vmla.gather %arg0([[RS]] : !shapex.ranked_shape<[4]>), [[BL]]([[RS]] : !shapex.ranked_shape<[4]>), out [[BUF]]([[RS]] : !shapex.ranked_shape<[4]>) {batch_dims = 0 : i64, dim = 0 : i64} : f32
+  %sort = "mhlo.sort"(%arg0) ( {
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 0 : i64, is_stable = false} : (tensor<4xf32>) -> tensor<4xf32>
+
+  // CHECK: return [[BUF]] : !vmla.buffer
+  return %sort : tensor<4xf32>
+}
+
+
+// CHECK-LABEL: func @sort2D
+func @sort2D(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> attributes { sym_visibility = "private" } {
+  // CHECK-DAG: [[C64:%.+]] = constant 64 : index
+  // CHECK-DAG: [[RS:%.+]] = shapex.const_ranked_shape : !shapex.ranked_shape<[4,4]>
+  // CHECK-DAG: [[BL:%.+]] = vmla.buffer.alloc byte_length = [[C64]] : !vmla.buffer
+  // CHECK-DAG: vmla.sort %arg0([[RS]] : !shapex.ranked_shape<[4,4]>), out [[BL]] : f32
+  // CHECK-DAG: [[BUF:%.+]] = vmla.buffer.alloc byte_length = [[C64]] : !vmla.buffer
+  // CHECK-DAG: vmla.gather %arg0([[RS]] : !shapex.ranked_shape<[4,4]>), [[BL]]([[RS]] : !shapex.ranked_shape<[4,4]>), out [[BUF]]([[RS]] : !shapex.ranked_shape<[4,4]>) {batch_dims = 1 : i64, dim = 1 : i64} : f32
+  %sort = "mhlo.sort"(%arg0) ( {
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 1 : i64, is_stable = false} : (tensor<4x4xf32>) -> tensor<4x4xf32>
+
+  // CHECK: return [[BUF]] : !vmla.buffer
+  return %sort : tensor<4x4xf32>
+}
diff --git a/iree/compiler/Dialect/VMLA/Conversion/VMLAToVM/ConvertVMLAToVM.cpp b/iree/compiler/Dialect/VMLA/Conversion/VMLAToVM/ConvertVMLAToVM.cpp
index 7e4f6ea..1b66485 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/VMLAToVM/ConvertVMLAToVM.cpp
+++ b/iree/compiler/Dialect/VMLA/Conversion/VMLAToVM/ConvertVMLAToVM.cpp
@@ -321,6 +321,7 @@
   VMLA_TYPED_IMPORT_OP(IREE::VMLA::FloorOp, "vmla.floor");
   VMLA_TYPED_IMPORT_OP(IREE::VMLA::CeilOp, "vmla.ceil");
   VMLA_TYPED_IMPORT_OP(IREE::VMLA::RoundOp, "vmla.round");
+  VMLA_TYPED_IMPORT_OP(IREE::VMLA::SortOp, "vmla.sort");
 
   patterns.insert<VMLAConvertImportOpConversion>(context, importSymbols,
                                                  typeConverter, "vmla.convert");
diff --git a/iree/compiler/Dialect/VMLA/IR/VMLAOps.td b/iree/compiler/Dialect/VMLA/IR/VMLAOps.td
index f5cd0fe..422fed3 100644
--- a/iree/compiler/Dialect/VMLA/IR/VMLAOps.td
+++ b/iree/compiler/Dialect/VMLA/IR/VMLAOps.td
@@ -422,7 +422,7 @@
 }
 
 //===----------------------------------------------------------------------===//
-// VMLA Ops: Convultion
+// VMLA Ops: Convolution
 //===----------------------------------------------------------------------===//
 
 def VLMA_ConvOp : VMLA_Op<"conv", [VMLA_IncludeShapes]> {
@@ -460,6 +460,46 @@
 }
 
 //===----------------------------------------------------------------------===//
+// VMLA Ops: Sorting
+//===----------------------------------------------------------------------===//
+
+def VMLA_SortPseudoOp : VMLA_Op<"sort.pseudo"> {
+  let summary = "Tensor-level pseudo-op of VMLA::SortOp.";
+  let description = [{
+    This is a tensor-level version of VMLA::SortOp, to facilitate
+    the lowering process.
+
+    This operation generates a sorted index list along the last dimension,
+    performing batch-wise along all other dimensions.
+  }];
+  let arguments = (ins
+    AnyTensor:$value
+  );
+  let results = (outs
+    I32Tensor:$dst
+  );
+
+  let assemblyFormat = [{
+    $value attr-dict `:` `(`type($value)`)` `->` type($dst)
+  }];
+}
+
+def VMLA_SortOp : VMLA_ElementTypeOp<"sort", [VMLA_IncludeShapes]> {
+  let arguments = (ins
+    VMLA_Buffer:$src,
+    VMLA_Shape:$src_shape,
+    VMLA_Buffer:$dst,
+    VMLA_AnyTypeAttr:$element_type
+  );
+
+  let assemblyFormat = [{
+    $src`(`$src_shape `:` type($src_shape)`)``,`
+    `out` $dst attr-dict `:` $element_type
+  }];
+}
+
+
+//===----------------------------------------------------------------------===//
 // VMLA Ops: GEMM/GEMV
 //===----------------------------------------------------------------------===//
 
diff --git a/iree/compiler/Dialect/VMLA/Transforms/PreConversionLowering.cpp b/iree/compiler/Dialect/VMLA/Transforms/PreConversionLowering.cpp
index d937026..62dd921 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/PreConversionLowering.cpp
+++ b/iree/compiler/Dialect/VMLA/Transforms/PreConversionLowering.cpp
@@ -277,6 +277,96 @@
   }
 };
 
+// Lower mhlo::SortOp to an pseudo SortOp in the VMLA dialect. This
+// pseudo op generates a set of ordered indices for that array along the last
+// dimension. Then using a torch_index_select the values can be reordered to
+// support arbitrary inputs.
+//
+// TODO(suderman): This lowering only covers the case of ascending values, we
+// should support a separate descending value case by having separate
+// SortAscending and SortDescending operations.
+class LowerSortOp : public OpRewritePattern<mhlo::SortOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(mhlo::SortOp op,
+                                PatternRewriter &rewriter) const override {
+    auto operandTy = op.getOperand(0).getType().cast<RankedTensorType>();
+    bool lastDimension =
+        (op.dimension() == -1) || (op.dimension() == (operandTy.getRank() - 1));
+
+    // TODO(suderman): Add transpose to sort along the last dimension.
+    if (!lastDimension) return failure();
+
+    auto &comparator = op.comparator();
+    auto &block = comparator.getBlocks().front();
+    auto &operations = block.getOperations();
+    auto comparison = dyn_cast_or_null<mhlo::CompareOp>(&operations.front());
+
+    // First verify that the block is purely a return of a comparison. This
+    // handles sorting a single tensor of values.
+    if (!comparison) return failure();
+
+    auto returnOp =
+        dyn_cast_or_null<mhlo::ReturnOp>(&(*(++operations.begin())));
+    if (!returnOp) return failure();
+
+    if (returnOp.getOperand(0) != comparison.getResult()) return failure();
+
+    // Determine which operands being compared.
+    auto lhs = comparison.getOperand(0);
+    auto rhs = comparison.getOperand(1);
+    auto lhsIndex = -1;
+    auto rhsIndex = -1;
+    for (auto arg : llvm::enumerate(block.getArguments())) {
+      if (arg.value() == lhs) lhsIndex = arg.index();
+      if (arg.value() == rhs) rhsIndex = arg.index();
+    }
+
+    // This should never happen but best to check.
+    if (lhsIndex == -1) return failure();
+    if (rhsIndex == -1) return failure();
+
+    // They should not be the same.
+    if (lhsIndex == rhsIndex) return failure();
+
+    // Comparisons need to pull from same Sort operand..
+    auto lhsOperand = lhsIndex / 2;
+    auto rhsOperand = rhsIndex / 2;
+    if (lhsOperand != rhsOperand) return failure();
+
+    // Must be GT, GE, LT, or LE.
+    auto isGt = comparison.comparison_direction() == "GT" ||
+                comparison.comparison_direction() == "GE";
+    auto isLt = comparison.comparison_direction() == "LT" ||
+                comparison.comparison_direction() == "LE";
+    if (!isGt && !isLt) return failure();
+
+    bool operandParity = lhsIndex > rhsIndex;
+    auto isAscending = operandParity ^ isGt;
+    // TODO(suderman): Add support for descended sorting.
+    if (!isAscending) return failure();
+
+    auto operand = op.getOperand(lhsOperand);
+    auto sortedIndices = rewriter.create<VMLA::SortPseudoOp>(
+        op.getLoc(),
+        RankedTensorType::get(operandTy.getShape(), rewriter.getI32Type()),
+        operand);
+
+    llvm::SmallVector<Value, 6> sortedResults;
+    for (auto operand : op.getOperands()) {
+      auto tensorTy = operand.getType().cast<RankedTensorType>();
+      auto gathered = rewriter.create<mhlo::TorchIndexSelectOp>(
+          op.getLoc(), tensorTy, operand, sortedIndices,
+          /**dim=*/operandTy.getRank() - 1,
+          /**batch_dims=*/operandTy.getRank() - 1);
+      sortedResults.push_back(gathered);
+    }
+
+    rewriter.replaceOp(op, sortedResults);
+    return success();
+  }
+};
+
 class PreConversionLoweringPass
     : public PassWrapper<PreConversionLoweringPass, OperationPass<FuncOp>> {
  public:
@@ -310,6 +400,8 @@
     patterns.insert<LowerBroadcastInDimOp>(context);
     target.addIllegalOp<mhlo::BroadcastOp>();
     patterns.insert<LowerBroadcastOp>(context);
+    target.addIllegalOp<mhlo::SortOp>();
+    patterns.insert<LowerSortOp>(context);
 
     if (failed(applyPartialConversion(getOperation(), target, patterns))) {
       return signalPassFailure();
diff --git a/iree/compiler/Dialect/VMLA/Transforms/test/pre_conversion_lowering.mlir b/iree/compiler/Dialect/VMLA/Transforms/test/pre_conversion_lowering.mlir
index 3473e44..0b9cd82 100644
--- a/iree/compiler/Dialect/VMLA/Transforms/test/pre_conversion_lowering.mlir
+++ b/iree/compiler/Dialect/VMLA/Transforms/test/pre_conversion_lowering.mlir
@@ -17,6 +17,38 @@
 // -----
 
 // CHECK-LABEL: func @f
+func @f(%arg0 : tensor<4xf32>) -> tensor<4xf32> attributes { sym_visibility = "private" } {
+  // CHECK-DAG: [[SORT:%.+]] = vmla.sort.pseudo %arg0
+  // CHECK-DAG: [[GATHER:%.+]] = "mhlo.torch_index_select"(%arg0, [[SORT]]) {batch_dims = 0 : i64, dim = 0 : i64}
+  %sort = "mhlo.sort"(%arg0) ( {
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 0 : i64, is_stable = false} : (tensor<4xf32>) -> tensor<4xf32>
+
+  // CHECK: return [[GATHER]]
+  return %sort : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @f
+func @f(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> attributes { sym_visibility = "private" } {
+  // CHECK-DAG: [[SORT:%.+]] = vmla.sort.pseudo %arg0
+  // CHECK-DAG: [[GATHER:%.+]] = "mhlo.torch_index_select"(%arg0, [[SORT]]) {batch_dims = 1 : i64, dim = 1 : i64}
+  %sort = "mhlo.sort"(%arg0) ( {
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 1 : i64, is_stable = false} : (tensor<4x4xf32>) -> tensor<4x4xf32>
+
+  // CHECK return [[GATHER]]
+  return %sort : tensor<4x4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @f
 func @f(%arg0: tensor<3xf32>) -> tensor<4x3xf32> {
   // CHECK: "shapex.ranked_broadcast_in_dim"(%arg0, %rs4_3)
   %0 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<[1]> : tensor<1xi64>} : (tensor<3xf32>) -> tensor<4x3xf32>
diff --git a/iree/compiler/Dialect/VMLA/vmla.imports.mlir b/iree/compiler/Dialect/VMLA/vmla.imports.mlir
index 17d7e85..ff575b0 100644
--- a/iree/compiler/Dialect/VMLA/vmla.imports.mlir
+++ b/iree/compiler/Dialect/VMLA/vmla.imports.mlir
@@ -333,6 +333,20 @@
 vm.import @ceil.f32(%src : !vm.ref<!vmla.buffer>, %dst : !vm.ref<!vmla.buffer>)
 vm.import @round.f32(%src : !vm.ref<!vmla.buffer>, %dst : !vm.ref<!vmla.buffer>)
 
+
+vm.import @sort.i8(
+  %src : !vm.ref<!vmla.buffer>, %src_shape : i32 ...,
+  %dst : !vm.ref<!vmla.buffer>)
+vm.import @sort.i16(
+  %src : !vm.ref<!vmla.buffer>, %src_shape : i32 ...,
+  %dst : !vm.ref<!vmla.buffer>)
+vm.import @sort.i32(
+  %src : !vm.ref<!vmla.buffer>, %src_shape : i32 ...,
+  %dst : !vm.ref<!vmla.buffer>)
+vm.import @sort.f32(
+  %src : !vm.ref<!vmla.buffer>, %src_shape : i32 ...,
+  %dst : !vm.ref<!vmla.buffer>)
+
 //===----------------------------------------------------------------------===//
 // VMLA Ops: conversion
 //===----------------------------------------------------------------------===//
diff --git a/iree/hal/vmla/op_kernels.h b/iree/hal/vmla/op_kernels.h
index c6ab9d6..ba5b8bc 100644
--- a/iree/hal/vmla/op_kernels.h
+++ b/iree/hal/vmla/op_kernels.h
@@ -168,6 +168,12 @@
                         absl::Span<const int32_t> dimensions);
 };
 
+struct Sort {
+  template <typename T>
+  static Status Execute(absl::Span<const T> src_buffer,
+                        absl::Span<int32_t> dst_buffer, ShapeSpan src_shape);
+};
+
 struct Broadcast {
   template <typename T>
   static Status Execute(absl::Span<const T> src_buffer,
diff --git a/iree/hal/vmla/op_kernels_generic.h b/iree/hal/vmla/op_kernels_generic.h
index d3545d1..0b4f904 100644
--- a/iree/hal/vmla/op_kernels_generic.h
+++ b/iree/hal/vmla/op_kernels_generic.h
@@ -15,7 +15,10 @@
 #ifndef IREE_HAL_VMLA_OP_KERNELS_GENERIC_H_
 #define IREE_HAL_VMLA_OP_KERNELS_GENERIC_H_
 
+#include <algorithm>
 #include <cmath>
+#include <iostream>
+#include <numeric>
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/inlined_vector.h"
@@ -519,6 +522,25 @@
 }
 
 template <typename T>
+Status Sort::Execute(absl::Span<const T> src_buffer,
+                     absl::Span<int32_t> dst_buffer, ShapeSpan src_shape) {
+  int elements = src_buffer.size();
+  const int sort_size = src_shape.back();
+
+  for (int i = 0; i < elements; i += sort_size) {
+    auto src_subspan = src_buffer.subspan(i, sort_size);
+    auto dst_subspan = dst_buffer.subspan(i, sort_size);
+    std::iota(dst_subspan.begin(), dst_subspan.end(), 0);
+    std::stable_sort(dst_subspan.begin(), dst_subspan.end(),
+                     [&src_subspan](int32_t i1, int32_t i2) {
+                       return src_subspan[i1] < src_subspan[i2];
+                     });
+  }
+
+  return OkStatus();
+}
+
+template <typename T>
 Status Broadcast::Execute(absl::Span<const T> src_buffer,
                           absl::Span<T> dst_buffer) {
   for (size_t i = 0; i < dst_buffer.size(); ++i) {
diff --git a/iree/hal/vmla/vmla_module.cc b/iree/hal/vmla/vmla_module.cc
index 09dbb3c..5852de0 100644
--- a/iree/hal/vmla/vmla_module.cc
+++ b/iree/hal/vmla/vmla_module.cc
@@ -642,6 +642,19 @@
   IREE_VMLA_UNARY_OP(CeilF32, kernels::Ceil, float);
   IREE_VMLA_UNARY_OP(RoundF32, kernels::Round, float);
 
+#define IREE_VMLA_SORT_OP(name, type)                                        \
+  Status name(const vm::ref<Buffer>& src, iree_vmla_shape_t src_shape,       \
+              const vm::ref<Buffer>& dst) {                                  \
+    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                            \
+    return kernels::Sort::Execute<type>(src->As<type>(), dst->As<int32_t>(), \
+                                        src_shape);                          \
+  }
+
+  IREE_VMLA_SORT_OP(SortI8, int8_t);
+  IREE_VMLA_SORT_OP(SortI16, int16_t);
+  IREE_VMLA_SORT_OP(SortI32, int32_t);
+  IREE_VMLA_SORT_OP(SortF32, float);
+
   //===--------------------------------------------------------------------===//
   // VMLA Ops: conversion
   //===--------------------------------------------------------------------===//
@@ -970,6 +983,10 @@
     vm::MakeNativeFunction("floor.f32", &VMLAModuleState::FloorF32),
     vm::MakeNativeFunction("ceil.f32", &VMLAModuleState::CeilF32),
     vm::MakeNativeFunction("round.f32", &VMLAModuleState::RoundF32),
+    vm::MakeNativeFunction("sort.i8", &VMLAModuleState::SortI8),
+    vm::MakeNativeFunction("sort.i16", &VMLAModuleState::SortI16),
+    vm::MakeNativeFunction("sort.i32", &VMLAModuleState::SortI32),
+    vm::MakeNativeFunction("sort.f32", &VMLAModuleState::SortF32),
     vm::MakeNativeFunction("finite.f32", &VMLAModuleState::FiniteF32),
 
     vm::MakeNativeFunction("convert.i8.i16", &VMLAModuleState::ConvertI8I16),
diff --git a/iree/test/e2e/xla_ops/sort.mlir b/iree/test/e2e/xla_ops/sort.mlir
new file mode 100644
index 0000000..1820d8d
--- /dev/null
+++ b/iree/test/e2e/xla_ops/sort.mlir
@@ -0,0 +1,40 @@
+func @sort1D() attributes { iree.module.export } {
+  %input = iree.unfoldable_constant dense<[3, 2, 1, 4]> : tensor<4xi32>
+
+  %sort = "mhlo.sort"(%input) ( {
+  ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 0 : i64, is_stable = false} : (tensor<4xi32>) -> tensor<4xi32>
+
+  check.expect_eq_const(%sort, dense<[1, 2, 3, 4]> : tensor<4xi32>) : tensor<4xi32>
+  return
+}
+
+func @sort2D() attributes { iree.module.export } {
+  %input = iree.unfoldable_constant dense<[[1, 2, 3, 4],
+                                           [4, 3, 2, 1]]> : tensor<2x4xi32>
+
+  %sort = "mhlo.sort"(%input) ( {
+  ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 1 : i64, is_stable = false} : (tensor<2x4xi32>) -> tensor<2x4xi32>
+
+  check.expect_eq_const(%sort, dense<[[1, 2, 3, 4], [1, 2, 3, 4]]> : tensor<2x4xi32>) : tensor<2x4xi32>
+  return
+}
+
+func @sort3D() attributes { iree.module.export } {
+  %input = iree.unfoldable_constant dense<[[[1, 2, 3, 4],
+                                            [4, 3, 2, 1]]]> : tensor<1x2x4xi32>
+
+  %sort = "mhlo.sort"(%input) ( {
+  ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>):  // no predecessors
+    %compare = "mhlo.compare"(%arg1, %arg2) {comparison_direction = "GT"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+    "mhlo.return"(%compare) : (tensor<i1>) -> ()
+  }) {dimension = 2 : i64, is_stable = false} : (tensor<1x2x4xi32>) -> tensor<1x2x4xi32>
+
+  check.expect_eq_const(%sort, dense<[[[1, 2, 3, 4], [1, 2, 3, 4]]]> : tensor<1x2x4xi32>) : tensor<1x2x4xi32>
+  return
+}
diff --git a/scripts/git/google_to_main.sh b/scripts/git/google_to_main.sh
index 8b8543d..b00fd9a 100755
--- a/scripts/git/google_to_main.sh
+++ b/scripts/git/google_to_main.sh
@@ -56,4 +56,7 @@
   echo "${BODY?}"
   exit 1
 fi
-gh pr create --base main --title="${TITLE?}" --body="${BODY?}"
+
+# Workaround https://github.com/cli/cli/issues/1820
+GITHUB_USERNAME="$(gh config get -h github.com user)"
+gh pr create --base main --head="${GITHUB_USERNAME?}:${PR_BRANCH?}" --title="${TITLE?}" --body="${BODY?}"
diff --git a/scripts/git/main_to_google.sh b/scripts/git/main_to_google.sh
index ee07573..b9e56ae 100755
--- a/scripts/git/main_to_google.sh
+++ b/scripts/git/main_to_google.sh
@@ -56,4 +56,7 @@
   echo "${BODY?}"
   exit 1
 fi
-gh pr create --base google --title="${TITLE?}" --body="${BODY?}"
+
+# Workaround https://github.com/cli/cli/issues/1820
+GITHUB_USERNAME="$(gh config get -h github.com user)"
+gh pr create --base google --head="${GITHUB_USERNAME?}:${PR_BRANCH?}" --title="${TITLE?}" --body="${BODY?}"
diff --git a/scripts/git/update_tf_submodule.sh b/scripts/git/update_tf_submodule.sh
index 0031cfa..28799d1 100755
--- a/scripts/git/update_tf_submodule.sh
+++ b/scripts/git/update_tf_submodule.sh
@@ -72,4 +72,7 @@
   echo "${BODY?}"
   exit 1
 fi
-gh pr create --title="${TITLE?}" --body="${BODY?}" --base="${BASE_BRANCH?}"
+
+# Workaround https://github.com/cli/cli/issues/1820
+GITHUB_USERNAME="$(gh config get -h github.com user)"
+gh pr create --base="${BASE_BRANCH?}" --head="${GITHUB_USERNAME?}:${PR_BRANCH?}" --title="${TITLE?}" --body="${BODY?}"