From b17f3c80dce555dab2f92fad459097d7f8c028b5 Mon Sep 17 00:00:00 2001
From: Kaxil Naik <kaxilnaik@gmail.com>
Date: Tue, 11 Mar 2025 23:35:53 +0530
Subject: [PATCH] Move `chain`, `chain_linear` & `cross_downstream` to Task SDK

This functions are for DAG Authors to define relationship between multiple tasks in batch
---
 .../example_asset_with_watchers.py            |   3 +-
 .../example_dags/example_bash_decorator.py    |   2 +-
 airflow/example_dags/example_complex.py       |   2 +-
 .../example_short_circuit_decorator.py        |   2 +-
 .../example_short_circuit_operator.py         |   2 +-
 airflow/models/baseoperator.py                | 273 +-----------------
 airflow/utils/edgemodifier.py                 |  16 +-
 dev/perf/dags/elastic_dag.py                  |   2 +-
 docs/apache-airflow/core-concepts/dags.rst    |   6 +-
 newsfragments/aip-72.significant.rst          |  34 ++-
 .../performance_dag/performance_dag.py        |   2 +-
 task-sdk/src/airflow/sdk/__init__.py          |   8 +-
 .../airflow/sdk/definitions/baseoperator.py   | 264 +++++++++++++++++
 .../src/airflow/sdk/definitions/xcom_arg.py   |   2 +-
 .../task_sdk/definitions/test_baseoperator.py | 199 ++++++++++++-
 tests/models/test_baseoperator.py             | 193 +------------
 tests/system/example_empty.py                 |   2 +-
 tests/utils/test_task_group.py                |   2 +-
 18 files changed, 523 insertions(+), 491 deletions(-)

diff --git a/airflow/example_dags/example_asset_with_watchers.py b/airflow/example_dags/example_asset_with_watchers.py
index 4f65a7f93055b..aa430d57b0699 100644
--- a/airflow/example_dags/example_asset_with_watchers.py
+++ b/airflow/example_dags/example_asset_with_watchers.py
@@ -21,10 +21,9 @@
 from __future__ import annotations
 
 from airflow.decorators import task
-from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
 from airflow.providers.standard.triggers.file import FileDeleteTrigger
-from airflow.sdk import Asset, AssetWatcher
+from airflow.sdk import Asset, AssetWatcher, chain
 
 file_path = "/tmp/test"
 
diff --git a/airflow/example_dags/example_bash_decorator.py b/airflow/example_dags/example_bash_decorator.py
index c976853fffc4e..7c34065fb1ce9 100644
--- a/airflow/example_dags/example_bash_decorator.py
+++ b/airflow/example_dags/example_bash_decorator.py
@@ -21,8 +21,8 @@
 
 from airflow.decorators import dag, task
 from airflow.exceptions import AirflowSkipException
-from airflow.models.baseoperator import chain
 from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.sdk import chain
 from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.weekday import WeekDay
 
diff --git a/airflow/example_dags/example_complex.py b/airflow/example_dags/example_complex.py
index 6d7d504f13d51..601e5b4c79007 100644
--- a/airflow/example_dags/example_complex.py
+++ b/airflow/example_dags/example_complex.py
@@ -23,9 +23,9 @@
 
 import pendulum
 
-from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
 from airflow.providers.standard.operators.bash import BashOperator
+from airflow.sdk import chain
 
 with DAG(
     dag_id="example_complex",
diff --git a/airflow/example_dags/example_short_circuit_decorator.py b/airflow/example_dags/example_short_circuit_decorator.py
index d87946d94687b..53cb0f31cca5e 100644
--- a/airflow/example_dags/example_short_circuit_decorator.py
+++ b/airflow/example_dags/example_short_circuit_decorator.py
@@ -21,8 +21,8 @@
 import pendulum
 
 from airflow.decorators import dag, task
-from airflow.models.baseoperator import chain
 from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.sdk import chain
 from airflow.utils.trigger_rule import TriggerRule
 
 
diff --git a/airflow/example_dags/example_short_circuit_operator.py b/airflow/example_dags/example_short_circuit_operator.py
index eb632185aa361..b20f8acb8d4db 100644
--- a/airflow/example_dags/example_short_circuit_operator.py
+++ b/airflow/example_dags/example_short_circuit_operator.py
@@ -21,10 +21,10 @@
 
 import pendulum
 
-from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.providers.standard.operators.python import ShortCircuitOperator
+from airflow.sdk import chain
 from airflow.utils.trigger_rule import TriggerRule
 
 with DAG(
diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index 3a663492c1f12..a37c28ef6a335 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -26,7 +26,7 @@
 import functools
 import logging
 import operator
-from collections.abc import Collection, Iterable, Iterator, Sequence
+from collections.abc import Collection, Iterable, Iterator
 from datetime import datetime, timedelta
 from functools import singledispatchmethod
 from types import FunctionType
@@ -54,14 +54,16 @@
     NotMapped,
 )
 from airflow.models.taskinstance import TaskInstance, clear_task_instances
-from airflow.models.taskmixin import DependencyMixin
 from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator as TaskSDKAbstractOperator
 from airflow.sdk.definitions.baseoperator import (
-    get_merged_defaults as get_merged_defaults,  # Re-export for compat
+    # Re-export for compat
+    chain as chain,
+    chain_linear as chain_linear,
+    cross_downstream as cross_downstream,
+    get_merged_defaults as get_merged_defaults,
 )
 from airflow.sdk.definitions.context import Context
 from airflow.sdk.definitions.dag import BaseOperator as TaskSDKBaseOperator
-from airflow.sdk.definitions.edges import EdgeModifier as TaskSDKEdgeModifier
 from airflow.sdk.definitions.mappedoperator import MappedOperator
 from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
 from airflow.serialization.enums import DagAttributeTypes
@@ -72,7 +74,6 @@
 from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep
 from airflow.utils import timezone
 from airflow.utils.context import context_get_outlet_events
-from airflow.utils.edgemodifier import EdgeModifier
 from airflow.utils.operator_helpers import ExecutionCallableRunner
 from airflow.utils.operator_resources import Resources
 from airflow.utils.session import NEW_SESSION, provide_session
@@ -811,265 +812,3 @@ def iter_mapped_task_group_lengths(group) -> Iterator[int]:
                     group = group.parent_group
 
             return functools.reduce(operator.mul, iter_mapped_task_group_lengths(group))
-
-
-def chain(*tasks: DependencyMixin | Sequence[DependencyMixin]) -> None:
-    r"""
-    Given a number of tasks, builds a dependency chain.
-
-    This function accepts values of BaseOperator (aka tasks), EdgeModifiers (aka Labels), XComArg, TaskGroups,
-    or lists containing any mix of these types (or a mix in the same list). If you want to chain between two
-    lists you must ensure they have the same length.
-
-    Using classic operators/sensors:
-
-    .. code-block:: python
-
-        chain(t1, [t2, t3], [t4, t5], t6)
-
-    is equivalent to::
-
-          / -> t2 -> t4 \
-        t1               -> t6
-          \ -> t3 -> t5 /
-
-    .. code-block:: python
-
-        t1.set_downstream(t2)
-        t1.set_downstream(t3)
-        t2.set_downstream(t4)
-        t3.set_downstream(t5)
-        t4.set_downstream(t6)
-        t5.set_downstream(t6)
-
-    Using task-decorated functions aka XComArgs:
-
-    .. code-block:: python
-
-        chain(x1(), [x2(), x3()], [x4(), x5()], x6())
-
-    is equivalent to::
-
-          / -> x2 -> x4 \
-        x1               -> x6
-          \ -> x3 -> x5 /
-
-    .. code-block:: python
-
-        x1 = x1()
-        x2 = x2()
-        x3 = x3()
-        x4 = x4()
-        x5 = x5()
-        x6 = x6()
-        x1.set_downstream(x2)
-        x1.set_downstream(x3)
-        x2.set_downstream(x4)
-        x3.set_downstream(x5)
-        x4.set_downstream(x6)
-        x5.set_downstream(x6)
-
-    Using TaskGroups:
-
-    .. code-block:: python
-
-        chain(t1, task_group1, task_group2, t2)
-
-        t1.set_downstream(task_group1)
-        task_group1.set_downstream(task_group2)
-        task_group2.set_downstream(t2)
-
-
-    It is also possible to mix between classic operator/sensor, EdgeModifiers, XComArg, and TaskGroups:
-
-    .. code-block:: python
-
-        chain(t1, [Label("branch one"), Label("branch two")], [x1(), x2()], task_group1, x3())
-
-    is equivalent to::
-
-          / "branch one" -> x1 \
-        t1                      -> task_group1 -> x3
-          \ "branch two" -> x2 /
-
-    .. code-block:: python
-
-        x1 = x1()
-        x2 = x2()
-        x3 = x3()
-        label1 = Label("branch one")
-        label2 = Label("branch two")
-        t1.set_downstream(label1)
-        label1.set_downstream(x1)
-        t2.set_downstream(label2)
-        label2.set_downstream(x2)
-        x1.set_downstream(task_group1)
-        x2.set_downstream(task_group1)
-        task_group1.set_downstream(x3)
-
-        # or
-
-        x1 = x1()
-        x2 = x2()
-        x3 = x3()
-        t1.set_downstream(x1, edge_modifier=Label("branch one"))
-        t1.set_downstream(x2, edge_modifier=Label("branch two"))
-        x1.set_downstream(task_group1)
-        x2.set_downstream(task_group1)
-        task_group1.set_downstream(x3)
-
-
-    :param tasks: Individual and/or list of tasks, EdgeModifiers, XComArgs, or TaskGroups to set dependencies
-    """
-    for up_task, down_task in zip(tasks, tasks[1:]):
-        if isinstance(up_task, DependencyMixin):
-            up_task.set_downstream(down_task)
-            continue
-        if isinstance(down_task, DependencyMixin):
-            down_task.set_upstream(up_task)
-            continue
-        if not isinstance(up_task, Sequence) or not isinstance(down_task, Sequence):
-            raise TypeError(f"Chain not supported between instances of {type(up_task)} and {type(down_task)}")
-        up_task_list = up_task
-        down_task_list = down_task
-        if len(up_task_list) != len(down_task_list):
-            raise AirflowException(
-                f"Chain not supported for different length Iterable. "
-                f"Got {len(up_task_list)} and {len(down_task_list)}."
-            )
-        for up_t, down_t in zip(up_task_list, down_task_list):
-            up_t.set_downstream(down_t)
-
-
-def cross_downstream(
-    from_tasks: Sequence[DependencyMixin],
-    to_tasks: DependencyMixin | Sequence[DependencyMixin],
-):
-    r"""
-    Set downstream dependencies for all tasks in from_tasks to all tasks in to_tasks.
-
-    Using classic operators/sensors:
-
-    .. code-block:: python
-
-        cross_downstream(from_tasks=[t1, t2, t3], to_tasks=[t4, t5, t6])
-
-    is equivalent to::
-
-        t1 ---> t4
-           \ /
-        t2 -X -> t5
-           / \
-        t3 ---> t6
-
-    .. code-block:: python
-
-        t1.set_downstream(t4)
-        t1.set_downstream(t5)
-        t1.set_downstream(t6)
-        t2.set_downstream(t4)
-        t2.set_downstream(t5)
-        t2.set_downstream(t6)
-        t3.set_downstream(t4)
-        t3.set_downstream(t5)
-        t3.set_downstream(t6)
-
-    Using task-decorated functions aka XComArgs:
-
-    .. code-block:: python
-
-        cross_downstream(from_tasks=[x1(), x2(), x3()], to_tasks=[x4(), x5(), x6()])
-
-    is equivalent to::
-
-        x1 ---> x4
-           \ /
-        x2 -X -> x5
-           / \
-        x3 ---> x6
-
-    .. code-block:: python
-
-        x1 = x1()
-        x2 = x2()
-        x3 = x3()
-        x4 = x4()
-        x5 = x5()
-        x6 = x6()
-        x1.set_downstream(x4)
-        x1.set_downstream(x5)
-        x1.set_downstream(x6)
-        x2.set_downstream(x4)
-        x2.set_downstream(x5)
-        x2.set_downstream(x6)
-        x3.set_downstream(x4)
-        x3.set_downstream(x5)
-        x3.set_downstream(x6)
-
-    It is also possible to mix between classic operator/sensor and XComArg tasks:
-
-    .. code-block:: python
-
-        cross_downstream(from_tasks=[t1, x2(), t3], to_tasks=[x1(), t2, x3()])
-
-    is equivalent to::
-
-        t1 ---> x1
-           \ /
-        x2 -X -> t2
-           / \
-        t3 ---> x3
-
-    .. code-block:: python
-
-        x1 = x1()
-        x2 = x2()
-        x3 = x3()
-        t1.set_downstream(x1)
-        t1.set_downstream(t2)
-        t1.set_downstream(x3)
-        x2.set_downstream(x1)
-        x2.set_downstream(t2)
-        x2.set_downstream(x3)
-        t3.set_downstream(x1)
-        t3.set_downstream(t2)
-        t3.set_downstream(x3)
-
-    :param from_tasks: List of tasks or XComArgs to start from.
-    :param to_tasks: List of tasks or XComArgs to set as downstream dependencies.
-    """
-    for task in from_tasks:
-        task.set_downstream(to_tasks)
-
-
-def chain_linear(*elements: DependencyMixin | Sequence[DependencyMixin]):
-    """
-    Simplify task dependency definition.
-
-    E.g.: suppose you want precedence like so::
-
-            ╭─op2─╮ ╭─op4─╮
-        op1─┤     ├─├─op5─┤─op7
-            ╰-op3─╯ ╰-op6─╯
-
-    Then you can accomplish like so::
-
-        chain_linear(op1, [op2, op3], [op4, op5, op6], op7)
-
-    :param elements: a list of operators / lists of operators
-    """
-    if not elements:
-        raise ValueError("No tasks provided; nothing to do.")
-    prev_elem = None
-    deps_set = False
-    for curr_elem in elements:
-        if isinstance(curr_elem, (EdgeModifier, TaskSDKEdgeModifier)):
-            raise ValueError("Labels are not supported by chain_linear")
-        if prev_elem is not None:
-            for task in prev_elem:
-                task >> curr_elem
-                if not deps_set:
-                    deps_set = True
-        prev_elem = [curr_elem] if isinstance(curr_elem, DependencyMixin) else curr_elem
-    if not deps_set:
-        raise ValueError("No dependencies were set. Did you forget to expand with `*`?")
diff --git a/airflow/utils/edgemodifier.py b/airflow/utils/edgemodifier.py
index 949511c3b4072..9e0d3da2f3326 100644
--- a/airflow/utils/edgemodifier.py
+++ b/airflow/utils/edgemodifier.py
@@ -16,17 +16,5 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
-import airflow.sdk
-
-if TYPE_CHECKING:
-    from airflow.typing_compat import TypeAlias
-
-EdgeModifier: TypeAlias = airflow.sdk.definitions.edges.EdgeModifier
-
-
-# Factory functions
-def Label(label: str):
-    """Create an EdgeModifier that sets a human-readable label on the edge."""
-    return EdgeModifier(label=label)
+# Re-export for compat
+from airflow.sdk.definitions.edges import Label as Label
diff --git a/dev/perf/dags/elastic_dag.py b/dev/perf/dags/elastic_dag.py
index 30bfc9acf992b..f05a638dabdc6 100644
--- a/dev/perf/dags/elastic_dag.py
+++ b/dev/perf/dags/elastic_dag.py
@@ -22,9 +22,9 @@
 from datetime import datetime, timedelta
 from enum import Enum
 
-from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
 from airflow.providers.standard.operators.bash import BashOperator
+from airflow.sdk import chain
 
 # DAG File used in performance tests. Its shape can be configured by environment variables.
 RE_TIME_DELTA = re.compile(
diff --git a/docs/apache-airflow/core-concepts/dags.rst b/docs/apache-airflow/core-concepts/dags.rst
index 1cf134bdc8dc3..b20c724c3e761 100644
--- a/docs/apache-airflow/core-concepts/dags.rst
+++ b/docs/apache-airflow/core-concepts/dags.rst
@@ -115,7 +115,7 @@ Or, you can also use the more explicit ``set_upstream`` and ``set_downstream`` m
 
 There are also shortcuts to declaring more complex dependencies. If you want to make a list of tasks depend on another list of tasks, you can't use either of the approaches above, so you need to use ``cross_downstream``::
 
-    from airflow.models.baseoperator import cross_downstream
+    from airflow.sdk import cross_downstream
 
     # Replaces
     # [op1, op2] >> op3
@@ -124,7 +124,7 @@ There are also shortcuts to declaring more complex dependencies. If you want to
 
 And if you want to chain together dependencies, you can use ``chain``::
 
-    from airflow.models.baseoperator import chain
+    from airflow.sdk import chain
 
     # Replaces op1 >> op2 >> op3 >> op4
     chain(op1, op2, op3, op4)
@@ -134,7 +134,7 @@ And if you want to chain together dependencies, you can use ``chain``::
 
 Chain can also do *pairwise* dependencies for lists the same size (this is different from the *cross dependencies* created by ``cross_downstream``!)::
 
-    from airflow.models.baseoperator import chain
+    from airflow.sdk import chain
 
     # Replaces
     # op1 >> op2 >> op4 >> op6
diff --git a/newsfragments/aip-72.significant.rst b/newsfragments/aip-72.significant.rst
index a20472372bfde..5b5b5114c032b 100644
--- a/newsfragments/aip-72.significant.rst
+++ b/newsfragments/aip-72.significant.rst
@@ -34,7 +34,7 @@ As part of this change the following breaking changes have occurred:
 
   .. code-block:: bash
 
-      airflow db drop-archived -t "_xcAm_archive"
+      airflow db drop-archived -t "_xcom_archive"
 
 - The ability to specify scheduling conditions for an operator via the ``deps`` class attribute has been removed.
 
@@ -46,7 +46,37 @@ As part of this change the following breaking changes have occurred:
 
   Any occurrences of imports from ``airflow.models.baseoperatorlink`` will need to be updated to ``airflow.sdk.definitions.baseoperatorlink``
 
-- With the We have removed DAG level settings that control the UI behaviour.
+- ``chain``, ``chain_linear`` and ``cross_downstream`` have been moved to the task SDK.
+
+  Any occurrences of imports from ``airflow.models.baseoperator`` will need to be updated to ``airflow.sdk``
+
+  Old imports:
+
+  .. code-block:: python
+
+      from airflow.models.baseoperator import chain, chain_linear, cross_downstream
+
+  New imports:
+
+  .. code-block:: python
+
+      from airflow.sdk import chain, chain_linear, cross_downstream
+
+- The ``Label`` class has been moved to the task SDK.
+
+  Old imports:
+
+  .. code-block:: python
+
+      from airflow.utils.edgemodifier import Label
+
+  New imports:
+
+  .. code-block:: python
+
+      from airflow.sdk import Label
+
+- We have removed DAG level settings that control the UI behaviour.
   These are now as per-user settings controlled by the UI
 
   - ``default_view``
diff --git a/performance/src/performance_dags/performance_dag/performance_dag.py b/performance/src/performance_dags/performance_dag/performance_dag.py
index c3e7a529a0dec..7286e2f6e75e2 100644
--- a/performance/src/performance_dags/performance_dag/performance_dag.py
+++ b/performance/src/performance_dags/performance_dag/performance_dag.py
@@ -54,9 +54,9 @@
 )
 
 from airflow import DAG
-from airflow.models.baseoperator import chain
 from airflow.operators.bash import BashOperator
 from airflow.operators.python import PythonOperator
+from airflow.sdk import chain
 from airflow.utils.trigger_rule import TriggerRule
 
 # DAG File used in performance tests. Its shape can be configured by environment variables.
diff --git a/task-sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py
index 5f4ecfd1d9007..f66c50fb3278d 100644
--- a/task-sdk/src/airflow/sdk/__init__.py
+++ b/task-sdk/src/airflow/sdk/__init__.py
@@ -38,6 +38,9 @@
     "Variable",
     "XComArg",
     "asset",
+    "chain",
+    "chain_linear",
+    "cross_downstream",
     "dag",
     "get_current_context",
     "get_parsing_context",
@@ -50,7 +53,7 @@
     from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAll, AssetAny, AssetWatcher
     from airflow.sdk.definitions.assets.decorators import asset
     from airflow.sdk.definitions.assets.metadata import Metadata
-    from airflow.sdk.definitions.baseoperator import BaseOperator
+    from airflow.sdk.definitions.baseoperator import BaseOperator, chain, chain_linear, cross_downstream
     from airflow.sdk.definitions.baseoperatorlink import BaseOperatorLink
     from airflow.sdk.definitions.connection import Connection
     from airflow.sdk.definitions.context import Context, get_current_context, get_parsing_context
@@ -81,6 +84,9 @@
     "Variable": ".definitions.variable",
     "XComArg": ".definitions.xcom_arg",
     "asset": ".definitions.asset.decorators",
+    "chain": ".definitions.baseoperator",
+    "chain_linear": ".definitions.baseoperator",
+    "cross_downstream": ".definitions.baseoperator",
     "dag": ".definitions.dag",
     "get_current_context": ".definitions.context",
     "get_parsing_context": ".definitions.context",
diff --git a/task-sdk/src/airflow/sdk/definitions/baseoperator.py b/task-sdk/src/airflow/sdk/definitions/baseoperator.py
index 097ba548bb800..b02eff27efb7a 100644
--- a/task-sdk/src/airflow/sdk/definitions/baseoperator.py
+++ b/task-sdk/src/airflow/sdk/definitions/baseoperator.py
@@ -50,11 +50,13 @@
     DEFAULT_WAIT_FOR_PAST_DEPENDS_BEFORE_SKIPPING,
     DEFAULT_WEIGHT_RULE,
     AbstractOperator,
+    DependencyMixin,
     TaskStateChangeCallback,
 )
 from airflow.sdk.definitions._internal.decorators import fixup_decorator_warning_stack
 from airflow.sdk.definitions._internal.node import validate_key
 from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, validate_instance_args
+from airflow.sdk.definitions.edges import EdgeModifier
 from airflow.sdk.definitions.mappedoperator import OperatorPartial, validate_mapping_kwargs
 from airflow.sdk.definitions.param import ParamsDict
 from airflow.task.priority_strategy import (
@@ -1577,3 +1579,265 @@ def resume_execution(self, next_method: str, next_kwargs: dict[str, Any] | None,
         # Grab the callable off the Operator/Task and add in any kwargs
         execute_callable = getattr(self, next_method)
         return execute_callable(context, **next_kwargs)
+
+
+def chain(*tasks: DependencyMixin | Sequence[DependencyMixin]) -> None:
+    r"""
+    Given a number of tasks, builds a dependency chain.
+
+    This function accepts values of BaseOperator (aka tasks), EdgeModifiers (aka Labels), XComArg, TaskGroups,
+    or lists containing any mix of these types (or a mix in the same list). If you want to chain between two
+    lists you must ensure they have the same length.
+
+    Using classic operators/sensors:
+
+    .. code-block:: python
+
+        chain(t1, [t2, t3], [t4, t5], t6)
+
+    is equivalent to::
+
+          / -> t2 -> t4 \
+        t1               -> t6
+          \ -> t3 -> t5 /
+
+    .. code-block:: python
+
+        t1.set_downstream(t2)
+        t1.set_downstream(t3)
+        t2.set_downstream(t4)
+        t3.set_downstream(t5)
+        t4.set_downstream(t6)
+        t5.set_downstream(t6)
+
+    Using task-decorated functions aka XComArgs:
+
+    .. code-block:: python
+
+        chain(x1(), [x2(), x3()], [x4(), x5()], x6())
+
+    is equivalent to::
+
+          / -> x2 -> x4 \
+        x1               -> x6
+          \ -> x3 -> x5 /
+
+    .. code-block:: python
+
+        x1 = x1()
+        x2 = x2()
+        x3 = x3()
+        x4 = x4()
+        x5 = x5()
+        x6 = x6()
+        x1.set_downstream(x2)
+        x1.set_downstream(x3)
+        x2.set_downstream(x4)
+        x3.set_downstream(x5)
+        x4.set_downstream(x6)
+        x5.set_downstream(x6)
+
+    Using TaskGroups:
+
+    .. code-block:: python
+
+        chain(t1, task_group1, task_group2, t2)
+
+        t1.set_downstream(task_group1)
+        task_group1.set_downstream(task_group2)
+        task_group2.set_downstream(t2)
+
+
+    It is also possible to mix between classic operator/sensor, EdgeModifiers, XComArg, and TaskGroups:
+
+    .. code-block:: python
+
+        chain(t1, [Label("branch one"), Label("branch two")], [x1(), x2()], task_group1, x3())
+
+    is equivalent to::
+
+          / "branch one" -> x1 \
+        t1                      -> task_group1 -> x3
+          \ "branch two" -> x2 /
+
+    .. code-block:: python
+
+        x1 = x1()
+        x2 = x2()
+        x3 = x3()
+        label1 = Label("branch one")
+        label2 = Label("branch two")
+        t1.set_downstream(label1)
+        label1.set_downstream(x1)
+        t2.set_downstream(label2)
+        label2.set_downstream(x2)
+        x1.set_downstream(task_group1)
+        x2.set_downstream(task_group1)
+        task_group1.set_downstream(x3)
+
+        # or
+
+        x1 = x1()
+        x2 = x2()
+        x3 = x3()
+        t1.set_downstream(x1, edge_modifier=Label("branch one"))
+        t1.set_downstream(x2, edge_modifier=Label("branch two"))
+        x1.set_downstream(task_group1)
+        x2.set_downstream(task_group1)
+        task_group1.set_downstream(x3)
+
+
+    :param tasks: Individual and/or list of tasks, EdgeModifiers, XComArgs, or TaskGroups to set dependencies
+    """
+    for up_task, down_task in zip(tasks, tasks[1:]):
+        if isinstance(up_task, DependencyMixin):
+            up_task.set_downstream(down_task)
+            continue
+        if isinstance(down_task, DependencyMixin):
+            down_task.set_upstream(up_task)
+            continue
+        if not isinstance(up_task, Sequence) or not isinstance(down_task, Sequence):
+            raise TypeError(f"Chain not supported between instances of {type(up_task)} and {type(down_task)}")
+        up_task_list = up_task
+        down_task_list = down_task
+        if len(up_task_list) != len(down_task_list):
+            raise ValueError(
+                f"Chain not supported for different length Iterable. "
+                f"Got {len(up_task_list)} and {len(down_task_list)}."
+            )
+        for up_t, down_t in zip(up_task_list, down_task_list):
+            up_t.set_downstream(down_t)
+
+
+def cross_downstream(
+    from_tasks: Sequence[DependencyMixin],
+    to_tasks: DependencyMixin | Sequence[DependencyMixin],
+):
+    r"""
+    Set downstream dependencies for all tasks in from_tasks to all tasks in to_tasks.
+
+    Using classic operators/sensors:
+
+    .. code-block:: python
+
+        cross_downstream(from_tasks=[t1, t2, t3], to_tasks=[t4, t5, t6])
+
+    is equivalent to::
+
+        t1 ---> t4
+           \ /
+        t2 -X -> t5
+           / \
+        t3 ---> t6
+
+    .. code-block:: python
+
+        t1.set_downstream(t4)
+        t1.set_downstream(t5)
+        t1.set_downstream(t6)
+        t2.set_downstream(t4)
+        t2.set_downstream(t5)
+        t2.set_downstream(t6)
+        t3.set_downstream(t4)
+        t3.set_downstream(t5)
+        t3.set_downstream(t6)
+
+    Using task-decorated functions aka XComArgs:
+
+    .. code-block:: python
+
+        cross_downstream(from_tasks=[x1(), x2(), x3()], to_tasks=[x4(), x5(), x6()])
+
+    is equivalent to::
+
+        x1 ---> x4
+           \ /
+        x2 -X -> x5
+           / \
+        x3 ---> x6
+
+    .. code-block:: python
+
+        x1 = x1()
+        x2 = x2()
+        x3 = x3()
+        x4 = x4()
+        x5 = x5()
+        x6 = x6()
+        x1.set_downstream(x4)
+        x1.set_downstream(x5)
+        x1.set_downstream(x6)
+        x2.set_downstream(x4)
+        x2.set_downstream(x5)
+        x2.set_downstream(x6)
+        x3.set_downstream(x4)
+        x3.set_downstream(x5)
+        x3.set_downstream(x6)
+
+    It is also possible to mix between classic operator/sensor and XComArg tasks:
+
+    .. code-block:: python
+
+        cross_downstream(from_tasks=[t1, x2(), t3], to_tasks=[x1(), t2, x3()])
+
+    is equivalent to::
+
+        t1 ---> x1
+           \ /
+        x2 -X -> t2
+           / \
+        t3 ---> x3
+
+    .. code-block:: python
+
+        x1 = x1()
+        x2 = x2()
+        x3 = x3()
+        t1.set_downstream(x1)
+        t1.set_downstream(t2)
+        t1.set_downstream(x3)
+        x2.set_downstream(x1)
+        x2.set_downstream(t2)
+        x2.set_downstream(x3)
+        t3.set_downstream(x1)
+        t3.set_downstream(t2)
+        t3.set_downstream(x3)
+
+    :param from_tasks: List of tasks or XComArgs to start from.
+    :param to_tasks: List of tasks or XComArgs to set as downstream dependencies.
+    """
+    for task in from_tasks:
+        task.set_downstream(to_tasks)
+
+
+def chain_linear(*elements: DependencyMixin | Sequence[DependencyMixin]):
+    """
+    Simplify task dependency definition.
+
+    E.g.: suppose you want precedence like so::
+
+            ╭─op2─╮ ╭─op4─╮
+        op1─┤     ├─├─op5─┤─op7
+            ╰-op3─╯ ╰-op6─╯
+
+    Then you can accomplish like so::
+
+        chain_linear(op1, [op2, op3], [op4, op5, op6], op7)
+
+    :param elements: a list of operators / lists of operators
+    """
+    if not elements:
+        raise ValueError("No tasks provided; nothing to do.")
+    prev_elem = None
+    deps_set = False
+    for curr_elem in elements:
+        if isinstance(curr_elem, EdgeModifier):
+            raise ValueError("Labels are not supported by chain_linear")
+        if prev_elem is not None:
+            for task in prev_elem:
+                task >> curr_elem
+                if not deps_set:
+                    deps_set = True
+        prev_elem = [curr_elem] if isinstance(curr_elem, DependencyMixin) else curr_elem
+    if not deps_set:
+        raise ValueError("No dependencies were set. Did you forget to expand with `*`?")
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index a090d0a6f9b86..012b1a9caad2f 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -35,8 +35,8 @@
 
 if TYPE_CHECKING:
     from airflow.sdk.definitions.baseoperator import BaseOperator
+    from airflow.sdk.definitions.edges import EdgeModifier
     from airflow.sdk.types import Operator
-    from airflow.utils.edgemodifier import EdgeModifier
 
 # Callable objects contained by MapXComArg. We only accept callables from
 # the user, but deserialize them into strings in a serialized XComArg for
diff --git a/task-sdk/tests/task_sdk/definitions/test_baseoperator.py b/task-sdk/tests/task_sdk/definitions/test_baseoperator.py
index 8c60b3fe7f298..0680b1df66a62 100644
--- a/task-sdk/tests/task_sdk/definitions/test_baseoperator.py
+++ b/task-sdk/tests/task_sdk/definitions/test_baseoperator.py
@@ -28,8 +28,18 @@
 import pytest
 import structlog
 
-from airflow.sdk.definitions.baseoperator import BaseOperator, BaseOperatorMeta, ExecutorSafeguard
+from airflow.decorators import task as task_decorator
+from airflow.sdk.definitions.baseoperator import (
+    BaseOperator,
+    BaseOperatorMeta,
+    ExecutorSafeguard,
+    chain,
+    chain_linear,
+    cross_downstream,
+)
 from airflow.sdk.definitions.dag import DAG
+from airflow.sdk.definitions.edges import Label
+from airflow.sdk.definitions.taskgroup import TaskGroup
 from airflow.sdk.definitions.template import literal
 from airflow.task.priority_strategy import _DownstreamPriorityWeightStrategy, _UpstreamPriorityWeightStrategy
 
@@ -264,6 +274,193 @@ def test_upstream_is_set_when_template_field_is_xcomarg(self):
         assert op1.task_id in op2.upstream_task_ids
         assert op2.task_id in op1.downstream_task_ids
 
+    def test_cross_downstream(self):
+        """Test if all dependencies between tasks are all set correctly."""
+        dag = DAG(dag_id="test_dag", schedule=None, start_date=datetime.now())
+        start_tasks = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 4)]
+        end_tasks = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(4, 7)]
+        cross_downstream(from_tasks=start_tasks, to_tasks=end_tasks)
+
+        for start_task in start_tasks:
+            assert set(start_task.get_direct_relatives(upstream=False)) == set(end_tasks)
+
+        # Begin test for `XComArgs`
+        xstart_tasks = [
+            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
+            for i in range(1, 4)
+        ]
+        xend_tasks = [
+            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
+            for i in range(4, 7)
+        ]
+        cross_downstream(from_tasks=xstart_tasks, to_tasks=xend_tasks)
+
+        for xstart_task in xstart_tasks:
+            assert set(xstart_task.operator.get_direct_relatives(upstream=False)) == {
+                xend_task.operator for xend_task in xend_tasks
+            }
+
+    def test_chain(self):
+        dag = DAG(dag_id="test_chain", schedule=None, start_date=datetime.now())
+
+        # Begin test for classic operators with `EdgeModifiers`
+        [label1, label2] = [Label(label=f"label{i}") for i in range(1, 3)]
+        [op1, op2, op3, op4, op5, op6] = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 7)]
+        chain(op1, [label1, label2], [op2, op3], [op4, op5], op6)
+
+        assert {op2, op3} == set(op1.get_direct_relatives(upstream=False))
+        assert [op4] == op2.get_direct_relatives(upstream=False)
+        assert [op5] == op3.get_direct_relatives(upstream=False)
+        assert {op4, op5} == set(op6.get_direct_relatives(upstream=True))
+
+        assert dag.get_edge_info(upstream_task_id=op1.task_id, downstream_task_id=op2.task_id) == {
+            "label": "label1"
+        }
+        assert dag.get_edge_info(upstream_task_id=op1.task_id, downstream_task_id=op3.task_id) == {
+            "label": "label2"
+        }
+
+        # Begin test for `XComArgs` with `EdgeModifiers`
+        [xlabel1, xlabel2] = [Label(label=f"xcomarg_label{i}") for i in range(1, 3)]
+        [xop1, xop2, xop3, xop4, xop5, xop6] = [
+            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
+            for i in range(1, 7)
+        ]
+        chain(xop1, [xlabel1, xlabel2], [xop2, xop3], [xop4, xop5], xop6)
+
+        assert {xop2.operator, xop3.operator} == set(xop1.operator.get_direct_relatives(upstream=False))
+        assert [xop4.operator] == xop2.operator.get_direct_relatives(upstream=False)
+        assert [xop5.operator] == xop3.operator.get_direct_relatives(upstream=False)
+        assert {xop4.operator, xop5.operator} == set(xop6.operator.get_direct_relatives(upstream=True))
+
+        assert dag.get_edge_info(
+            upstream_task_id=xop1.operator.task_id, downstream_task_id=xop2.operator.task_id
+        ) == {"label": "xcomarg_label1"}
+        assert dag.get_edge_info(
+            upstream_task_id=xop1.operator.task_id, downstream_task_id=xop3.operator.task_id
+        ) == {"label": "xcomarg_label2"}
+
+        # Begin test for `TaskGroups`
+        [tg1, tg2] = [TaskGroup(group_id=f"tg{i}", dag=dag) for i in range(1, 3)]
+        [op1, op2] = [BaseOperator(task_id=f"task{i}", dag=dag) for i in range(1, 3)]
+        [tgop1, tgop2] = [
+            BaseOperator(task_id=f"task_group_task{i}", task_group=tg1, dag=dag) for i in range(1, 3)
+        ]
+        [tgop3, tgop4] = [
+            BaseOperator(task_id=f"task_group_task{i}", task_group=tg2, dag=dag) for i in range(1, 3)
+        ]
+        chain(op1, tg1, tg2, op2)
+
+        assert {tgop1, tgop2} == set(op1.get_direct_relatives(upstream=False))
+        assert {tgop3, tgop4} == set(tgop1.get_direct_relatives(upstream=False))
+        assert {tgop3, tgop4} == set(tgop2.get_direct_relatives(upstream=False))
+        assert [op2] == tgop3.get_direct_relatives(upstream=False)
+        assert [op2] == tgop4.get_direct_relatives(upstream=False)
+
+    def test_chain_linear(self):
+        dag = DAG(dag_id="test_chain_linear", schedule=None, start_date=datetime.now())
+
+        t1, t2, t3, t4, t5, t6, t7 = (BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 8))
+        chain_linear(t1, [t2, t3, t4], [t5, t6], t7)
+
+        assert set(t1.get_direct_relatives(upstream=False)) == {t2, t3, t4}
+        assert set(t2.get_direct_relatives(upstream=False)) == {t5, t6}
+        assert set(t3.get_direct_relatives(upstream=False)) == {t5, t6}
+        assert set(t7.get_direct_relatives(upstream=True)) == {t5, t6}
+
+        t1, t2, t3, t4, t5, t6 = (
+            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
+            for i in range(1, 7)
+        )
+        chain_linear(t1, [t2, t3], [t4, t5], t6)
+
+        assert set(t1.operator.get_direct_relatives(upstream=False)) == {t2.operator, t3.operator}
+        assert set(t2.operator.get_direct_relatives(upstream=False)) == {t4.operator, t5.operator}
+        assert set(t3.operator.get_direct_relatives(upstream=False)) == {t4.operator, t5.operator}
+        assert set(t6.operator.get_direct_relatives(upstream=True)) == {t4.operator, t5.operator}
+
+        # Begin test for `TaskGroups`
+        tg1, tg2 = (TaskGroup(group_id=f"tg{i}", dag=dag) for i in range(1, 3))
+        op1, op2 = (BaseOperator(task_id=f"task{i}", dag=dag) for i in range(1, 3))
+        tgop1, tgop2 = (
+            BaseOperator(task_id=f"task_group_task{i}", task_group=tg1, dag=dag) for i in range(1, 3)
+        )
+        tgop3, tgop4 = (
+            BaseOperator(task_id=f"task_group_task{i}", task_group=tg2, dag=dag) for i in range(1, 3)
+        )
+        chain_linear(op1, tg1, tg2, op2)
+
+        assert set(op1.get_direct_relatives(upstream=False)) == {tgop1, tgop2}
+        assert set(tgop1.get_direct_relatives(upstream=False)) == {tgop3, tgop4}
+        assert set(tgop2.get_direct_relatives(upstream=False)) == {tgop3, tgop4}
+        assert set(tgop3.get_direct_relatives(upstream=False)) == {op2}
+        assert set(tgop4.get_direct_relatives(upstream=False)) == {op2}
+
+        t1, t2 = (BaseOperator(task_id=f"t-{i}", dag=dag) for i in range(1, 3))
+        with pytest.raises(ValueError, match="Labels are not supported"):
+            chain_linear(t1, Label("hi"), t2)
+
+        with pytest.raises(ValueError, match="nothing to do"):
+            chain_linear()
+
+        with pytest.raises(ValueError, match="Did you forget to expand"):
+            chain_linear(t1)
+
+    def test_chain_not_support_type(self):
+        dag = DAG(dag_id="test_chain", schedule=None, start_date=datetime.now())
+        [op1, op2] = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 3)]
+        with pytest.raises(TypeError):
+            chain([op1, op2], 1)
+
+        # Begin test for `XComArgs`
+        [xop1, xop2] = [
+            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
+            for i in range(1, 3)
+        ]
+
+        with pytest.raises(TypeError):
+            chain([xop1, xop2], 1)
+
+        # Begin test for `EdgeModifiers`
+        with pytest.raises(TypeError):
+            chain([Label("labe1"), Label("label2")], 1)
+
+        # Begin test for `TaskGroups`
+        [tg1, tg2] = [TaskGroup(group_id=f"tg{i}", dag=dag) for i in range(1, 3)]
+
+        with pytest.raises(TypeError):
+            chain([tg1, tg2], 1)
+
+    def test_chain_different_length_iterable(self):
+        dag = DAG(dag_id="test_chain", schedule=None, start_date=datetime.now())
+        [label1, label2] = [Label(label=f"label{i}") for i in range(1, 3)]
+        [op1, op2, op3, op4, op5] = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 6)]
+
+        with pytest.raises(ValueError):
+            chain([op1, op2], [op3, op4, op5])
+
+        with pytest.raises(ValueError):
+            chain([op1, op2, op3], [label1, label2])
+
+        # Begin test for `XComArgs` with `EdgeModifiers`
+        [label3, label4] = [Label(label=f"xcomarg_label{i}") for i in range(1, 3)]
+        [xop1, xop2, xop3, xop4, xop5] = [
+            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
+            for i in range(1, 6)
+        ]
+
+        with pytest.raises(ValueError):
+            chain([xop1, xop2], [xop3, xop4, xop5])
+
+        with pytest.raises(ValueError):
+            chain([xop1, xop2, xop3], [label1, label2])
+
+        # Begin test for `TaskGroups`
+        [tg1, tg2, tg3, tg4, tg5] = [TaskGroup(group_id=f"tg{i}", dag=dag) for i in range(1, 6)]
+
+        with pytest.raises(ValueError):
+            chain([tg1, tg2], [tg3, tg4, tg5])
+
     def test_set_xcomargs_dependencies_works_recursively(self):
         with DAG("xcomargs_test", schedule=None):
             op1 = BaseOperator(task_id="op1")
diff --git a/tests/models/test_baseoperator.py b/tests/models/test_baseoperator.py
index ee45fb02933fb..c3e7b3bd32a1e 100644
--- a/tests/models/test_baseoperator.py
+++ b/tests/models/test_baseoperator.py
@@ -25,12 +25,9 @@
 import pytest
 
 from airflow.decorators import task as task_decorator
-from airflow.exceptions import AirflowException, TaskDeferralTimeout
+from airflow.exceptions import TaskDeferralTimeout
 from airflow.models.baseoperator import (
     BaseOperator,
-    chain,
-    chain_linear,
-    cross_downstream,
 )
 from airflow.models.dag import DAG
 from airflow.models.dagrun import DagRun
@@ -38,7 +35,6 @@
 from airflow.models.trigger import TriggerFailureReason
 from airflow.providers.common.compat.lineage.entities import File
 from airflow.providers.common.sql.operators import sql
-from airflow.utils.edgemodifier import Label
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.types import DagRunType
@@ -91,89 +87,6 @@ def test_trigger_rule_validation(self):
             task_id="test_valid_trigger_rule", dag=non_fail_fast_dag, trigger_rule=TriggerRule.ALWAYS
         )
 
-    def test_cross_downstream(self):
-        """Test if all dependencies between tasks are all set correctly."""
-        dag = DAG(dag_id="test_dag", schedule=None, start_date=datetime.now())
-        start_tasks = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 4)]
-        end_tasks = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(4, 7)]
-        cross_downstream(from_tasks=start_tasks, to_tasks=end_tasks)
-
-        for start_task in start_tasks:
-            assert set(start_task.get_direct_relatives(upstream=False)) == set(end_tasks)
-
-        # Begin test for `XComArgs`
-        xstart_tasks = [
-            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
-            for i in range(1, 4)
-        ]
-        xend_tasks = [
-            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
-            for i in range(4, 7)
-        ]
-        cross_downstream(from_tasks=xstart_tasks, to_tasks=xend_tasks)
-
-        for xstart_task in xstart_tasks:
-            assert set(xstart_task.operator.get_direct_relatives(upstream=False)) == {
-                xend_task.operator for xend_task in xend_tasks
-            }
-
-    def test_chain(self):
-        dag = DAG(dag_id="test_chain", schedule=None, start_date=datetime.now())
-
-        # Begin test for classic operators with `EdgeModifiers`
-        [label1, label2] = [Label(label=f"label{i}") for i in range(1, 3)]
-        [op1, op2, op3, op4, op5, op6] = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 7)]
-        chain(op1, [label1, label2], [op2, op3], [op4, op5], op6)
-
-        assert {op2, op3} == set(op1.get_direct_relatives(upstream=False))
-        assert [op4] == op2.get_direct_relatives(upstream=False)
-        assert [op5] == op3.get_direct_relatives(upstream=False)
-        assert {op4, op5} == set(op6.get_direct_relatives(upstream=True))
-
-        assert dag.get_edge_info(upstream_task_id=op1.task_id, downstream_task_id=op2.task_id) == {
-            "label": "label1"
-        }
-        assert dag.get_edge_info(upstream_task_id=op1.task_id, downstream_task_id=op3.task_id) == {
-            "label": "label2"
-        }
-
-        # Begin test for `XComArgs` with `EdgeModifiers`
-        [xlabel1, xlabel2] = [Label(label=f"xcomarg_label{i}") for i in range(1, 3)]
-        [xop1, xop2, xop3, xop4, xop5, xop6] = [
-            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
-            for i in range(1, 7)
-        ]
-        chain(xop1, [xlabel1, xlabel2], [xop2, xop3], [xop4, xop5], xop6)
-
-        assert {xop2.operator, xop3.operator} == set(xop1.operator.get_direct_relatives(upstream=False))
-        assert [xop4.operator] == xop2.operator.get_direct_relatives(upstream=False)
-        assert [xop5.operator] == xop3.operator.get_direct_relatives(upstream=False)
-        assert {xop4.operator, xop5.operator} == set(xop6.operator.get_direct_relatives(upstream=True))
-
-        assert dag.get_edge_info(
-            upstream_task_id=xop1.operator.task_id, downstream_task_id=xop2.operator.task_id
-        ) == {"label": "xcomarg_label1"}
-        assert dag.get_edge_info(
-            upstream_task_id=xop1.operator.task_id, downstream_task_id=xop3.operator.task_id
-        ) == {"label": "xcomarg_label2"}
-
-        # Begin test for `TaskGroups`
-        [tg1, tg2] = [TaskGroup(group_id=f"tg{i}", dag=dag) for i in range(1, 3)]
-        [op1, op2] = [BaseOperator(task_id=f"task{i}", dag=dag) for i in range(1, 3)]
-        [tgop1, tgop2] = [
-            BaseOperator(task_id=f"task_group_task{i}", task_group=tg1, dag=dag) for i in range(1, 3)
-        ]
-        [tgop3, tgop4] = [
-            BaseOperator(task_id=f"task_group_task{i}", task_group=tg2, dag=dag) for i in range(1, 3)
-        ]
-        chain(op1, tg1, tg2, op2)
-
-        assert {tgop1, tgop2} == set(op1.get_direct_relatives(upstream=False))
-        assert {tgop3, tgop4} == set(tgop1.get_direct_relatives(upstream=False))
-        assert {tgop3, tgop4} == set(tgop2.get_direct_relatives(upstream=False))
-        assert [op2] == tgop3.get_direct_relatives(upstream=False)
-        assert [op2] == tgop4.get_direct_relatives(upstream=False)
-
     def test_baseoperator_raises_exception_when_task_id_plus_taskgroup_id_exceeds_250_chars(self):
         """Test exception is raised when operator task id + taskgroup id > 250 chars."""
         dag = DAG(dag_id="foo", schedule=None, start_date=datetime.now())
@@ -201,110 +114,6 @@ def test_baseoperator_with_task_id_less_than_250_chars(self):
         except Exception as e:
             pytest.fail(f"Exception raised: {e}")
 
-    def test_chain_linear(self):
-        dag = DAG(dag_id="test_chain_linear", schedule=None, start_date=datetime.now())
-
-        t1, t2, t3, t4, t5, t6, t7 = (BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 8))
-        chain_linear(t1, [t2, t3, t4], [t5, t6], t7)
-
-        assert set(t1.get_direct_relatives(upstream=False)) == {t2, t3, t4}
-        assert set(t2.get_direct_relatives(upstream=False)) == {t5, t6}
-        assert set(t3.get_direct_relatives(upstream=False)) == {t5, t6}
-        assert set(t7.get_direct_relatives(upstream=True)) == {t5, t6}
-
-        t1, t2, t3, t4, t5, t6 = (
-            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
-            for i in range(1, 7)
-        )
-        chain_linear(t1, [t2, t3], [t4, t5], t6)
-
-        assert set(t1.operator.get_direct_relatives(upstream=False)) == {t2.operator, t3.operator}
-        assert set(t2.operator.get_direct_relatives(upstream=False)) == {t4.operator, t5.operator}
-        assert set(t3.operator.get_direct_relatives(upstream=False)) == {t4.operator, t5.operator}
-        assert set(t6.operator.get_direct_relatives(upstream=True)) == {t4.operator, t5.operator}
-
-        # Begin test for `TaskGroups`
-        tg1, tg2 = (TaskGroup(group_id=f"tg{i}", dag=dag) for i in range(1, 3))
-        op1, op2 = (BaseOperator(task_id=f"task{i}", dag=dag) for i in range(1, 3))
-        tgop1, tgop2 = (
-            BaseOperator(task_id=f"task_group_task{i}", task_group=tg1, dag=dag) for i in range(1, 3)
-        )
-        tgop3, tgop4 = (
-            BaseOperator(task_id=f"task_group_task{i}", task_group=tg2, dag=dag) for i in range(1, 3)
-        )
-        chain_linear(op1, tg1, tg2, op2)
-
-        assert set(op1.get_direct_relatives(upstream=False)) == {tgop1, tgop2}
-        assert set(tgop1.get_direct_relatives(upstream=False)) == {tgop3, tgop4}
-        assert set(tgop2.get_direct_relatives(upstream=False)) == {tgop3, tgop4}
-        assert set(tgop3.get_direct_relatives(upstream=False)) == {op2}
-        assert set(tgop4.get_direct_relatives(upstream=False)) == {op2}
-
-        t1, t2 = (BaseOperator(task_id=f"t-{i}", dag=dag) for i in range(1, 3))
-        with pytest.raises(ValueError, match="Labels are not supported"):
-            chain_linear(t1, Label("hi"), t2)
-
-        with pytest.raises(ValueError, match="nothing to do"):
-            chain_linear()
-
-        with pytest.raises(ValueError, match="Did you forget to expand"):
-            chain_linear(t1)
-
-    def test_chain_not_support_type(self):
-        dag = DAG(dag_id="test_chain", schedule=None, start_date=datetime.now())
-        [op1, op2] = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 3)]
-        with pytest.raises(TypeError):
-            chain([op1, op2], 1)
-
-        # Begin test for `XComArgs`
-        [xop1, xop2] = [
-            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
-            for i in range(1, 3)
-        ]
-
-        with pytest.raises(TypeError):
-            chain([xop1, xop2], 1)
-
-        # Begin test for `EdgeModifiers`
-        with pytest.raises(TypeError):
-            chain([Label("labe1"), Label("label2")], 1)
-
-        # Begin test for `TaskGroups`
-        [tg1, tg2] = [TaskGroup(group_id=f"tg{i}", dag=dag) for i in range(1, 3)]
-
-        with pytest.raises(TypeError):
-            chain([tg1, tg2], 1)
-
-    def test_chain_different_length_iterable(self):
-        dag = DAG(dag_id="test_chain", schedule=None, start_date=datetime.now())
-        [label1, label2] = [Label(label=f"label{i}") for i in range(1, 3)]
-        [op1, op2, op3, op4, op5] = [BaseOperator(task_id=f"t{i}", dag=dag) for i in range(1, 6)]
-
-        with pytest.raises(AirflowException):
-            chain([op1, op2], [op3, op4, op5])
-
-        with pytest.raises(AirflowException):
-            chain([op1, op2, op3], [label1, label2])
-
-        # Begin test for `XComArgs` with `EdgeModifiers`
-        [label3, label4] = [Label(label=f"xcomarg_label{i}") for i in range(1, 3)]
-        [xop1, xop2, xop3, xop4, xop5] = [
-            task_decorator(task_id=f"xcomarg_task{i}", python_callable=lambda: None, dag=dag)()
-            for i in range(1, 6)
-        ]
-
-        with pytest.raises(AirflowException):
-            chain([xop1, xop2], [xop3, xop4, xop5])
-
-        with pytest.raises(AirflowException):
-            chain([xop1, xop2, xop3], [label1, label2])
-
-        # Begin test for `TaskGroups`
-        [tg1, tg2, tg3, tg4, tg5] = [TaskGroup(group_id=f"tg{i}", dag=dag) for i in range(1, 6)]
-
-        with pytest.raises(AirflowException):
-            chain([tg1, tg2], [tg3, tg4, tg5])
-
     def test_lineage_composition(self):
         """
         Test composition with lineage
diff --git a/tests/system/example_empty.py b/tests/system/example_empty.py
index ca81ea08d0be5..a452cd6382376 100644
--- a/tests/system/example_empty.py
+++ b/tests/system/example_empty.py
@@ -18,9 +18,9 @@
 
 from datetime import datetime
 
-from airflow.models.baseoperator import chain
 from airflow.models.dag import DAG
 from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.sdk import chain
 
 DAG_ID = "example_empty"
 
diff --git a/tests/utils/test_task_group.py b/tests/utils/test_task_group.py
index cfaeaf8288b61..082d886fbc08c 100644
--- a/tests/utils/test_task_group.py
+++ b/tests/utils/test_task_group.py
@@ -1458,7 +1458,7 @@ def test_add_to_another_group():
 
 
 def test_task_group_edge_modifier_chain():
-    from airflow.models.baseoperator import chain
+    from airflow.sdk import chain
     from airflow.utils.edgemodifier import Label
 
     with DAG(dag_id="test", schedule=None, start_date=pendulum.DateTime(2022, 5, 20)) as dag: