diff --git a/tfx/dsl/components/base/base_component.py b/tfx/dsl/components/base/base_component.py
index c263e743b7..11d67a00a0 100644
--- a/tfx/dsl/components/base/base_component.py
+++ b/tfx/dsl/components/base/base_component.py
@@ -25,6 +25,7 @@
 from tfx.types.system_executions import SystemExecution
 from tfx.utils import abc_utils
 from tfx.utils import doc_controls
+import typing_extensions
 
 from google.protobuf import message
 
@@ -150,8 +151,11 @@ def _validate_spec(self, spec):
            'got %s instead.') %
           (self.__class__, self.__class__.SPEC_CLASS, spec))
 
+  # TODO(kmonte): Update this to Self once we're on 3.11 everywhere.
   @doc_controls.do_not_doc_in_subclasses
-  def with_platform_config(self, config: message.Message) -> 'BaseComponent':
+  def with_platform_config(
+      self, config: message.Message
+  ) -> typing_extensions.Self:
     """Attaches a proto-form platform config to a component.
 
     The config will be a per-node platform-specific config.
diff --git a/tfx/dsl/components/base/base_node.py b/tfx/dsl/components/base/base_node.py
index cfb0b6a120..a737f2da3d 100644
--- a/tfx/dsl/components/base/base_node.py
+++ b/tfx/dsl/components/base/base_node.py
@@ -26,6 +26,7 @@
 from tfx.utils import doc_controls
 from tfx.utils import json_utils
 from tfx.utils import name_utils
+import typing_extensions
 
 
 def _abstract_property() -> Any:
@@ -128,8 +129,9 @@ def component_id(self) -> str:
   def id(self, id: str) -> None:  # pylint: disable=redefined-builtin
     self._id = id
 
+  # TODO(kmonte): Update this to Self once we're on 3.11 everywhere
   @doc_controls.do_not_doc_in_subclasses
-  def with_id(self, id: str) -> 'BaseNode':  # pylint: disable=redefined-builtin
+  def with_id(self, id: str) -> typing_extensions.Self:  # pylint: disable=redefined-builtin
     self._id = id
     return self
 
@@ -166,10 +168,10 @@ def node_execution_options(
   ):
     self._node_execution_options = copy.deepcopy(node_execution_options)
 
+  # TODO(kmonte): Update this to Self once we're on 3.11 everywhere
   def with_node_execution_options(
-      self,
-      node_execution_options: utils.NodeExecutionOptions
-  ) -> 'BaseNode':
+      self, node_execution_options: utils.NodeExecutionOptions
+  ) -> typing_extensions.Self:
     self.node_execution_options = node_execution_options
     return self