Replace _ArtiactMultimap aliases with typing_utils.ArtifactMultiMap.

PiperOrigin-RevId: 389763694

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

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):

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():