From 87a985d6722deab56aa6eddb61d51ba3a96135fb Mon Sep 17 00:00:00 2001
From: jjong <jjong@google.com>
Date: Mon, 9 Aug 2021 17:33:20 -0700
Subject: [PATCH] Replace _ArtiactMultimap aliases with
 typing_utils.ArtifactMultiMap.

PiperOrigin-RevId: 389763694
---
 tfx/dsl/input_resolution/resolver_function.py |  4 ++--
 tfx/dsl/input_resolution/resolver_op.py       | 11 ++++++-----
 tfx/orchestration/portable/inputs_utils.py    |  5 +----
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/tfx/dsl/input_resolution/resolver_function.py b/tfx/dsl/input_resolution/resolver_function.py
index 71078fda1d..ed2ce50bb4 100644
--- a/tfx/dsl/input_resolution/resolver_function.py
+++ b/tfx/dsl/input_resolution/resolver_function.py
@@ -22,8 +22,8 @@
 class ResolverFunction:
   """ResolverFunction represents a traceable function of resolver operators.
 
-  ResolverFunction as a whole, takes an ArtifactMultimap as an argument and
-  returns ArtifactMultimap.
+  ResolverFunction as a whole, takes an ArtifactMultiMap as an argument and
+  returns an ArtifactMultiMap.
 
   Usage:
       @resolver_function
diff --git a/tfx/dsl/input_resolution/resolver_op.py b/tfx/dsl/input_resolution/resolver_op.py
index 27967db6d7..ba6e827610 100644
--- a/tfx/dsl/input_resolution/resolver_op.py
+++ b/tfx/dsl/input_resolution/resolver_op.py
@@ -13,16 +13,14 @@
 # limitations under the License.
 """Module for ResolverOp and its related definitions."""
 import abc
-from typing import Any, ClassVar, Generic, Mapping, Sequence, Type, TypeVar, Union
+from typing import Any, ClassVar, Generic, Mapping, Type, TypeVar, Union
 
 import attr
-import tfx.types
 from tfx.utils import json_utils
+from tfx.utils import typing_utils
 
 import ml_metadata as mlmd
 
-ArtifactMultimap = Mapping[str, Sequence[tfx.types.Artifact]]
-
 
 # Mark frozen as context instance may be used across multiple operator
 # invocations.
@@ -220,7 +218,10 @@ class ResolverOp(metaclass=_ResolverOpMeta):
   """
 
   @abc.abstractmethod
-  def apply(self, input_dict: ArtifactMultimap) -> ArtifactMultimap:
+  def apply(
+      self,
+      input_dict: typing_utils.ArtifactMultiMap,
+  ) -> typing_utils.ArtifactMultiMap:
     """Implementation of the operator."""
 
   def set_context(self, context: Context):
diff --git a/tfx/orchestration/portable/inputs_utils.py b/tfx/orchestration/portable/inputs_utils.py
index b74eb7839d..1c05209a2c 100644
--- a/tfx/orchestration/portable/inputs_utils.py
+++ b/tfx/orchestration/portable/inputs_utils.py
@@ -16,7 +16,6 @@
 
 from absl import logging
 from tfx import types
-from tfx.dsl.input_resolution import resolver_op
 from tfx.orchestration import data_types_utils
 from tfx.orchestration import metadata
 from tfx.orchestration.portable.input_resolution import exceptions
@@ -31,8 +30,6 @@
 import ml_metadata as mlmd
 from ml_metadata.proto import metadata_store_pb2
 
-_ArtifactMultimap = resolver_op.ArtifactMultimap
-
 
 def get_qualified_artifacts(
     metadata_handler: metadata.Metadata,
@@ -123,7 +120,7 @@ def _resolve_single_channel(
 
 def _resolve_initial_dict(
     metadata_handler: metadata.Metadata,
-    node_inputs: pipeline_pb2.NodeInputs) -> _ArtifactMultimap:
+    node_inputs: pipeline_pb2.NodeInputs) -> typing_utils.ArtifactMultiMap:
   """Resolves initial input dict from input channel definition."""
   result = {}
   for key, input_spec in node_inputs.inputs.items():