From 90cd44503ed113b348b0db6506f2db9c7f2b8052 Mon Sep 17 00:00:00 2001
From: Advait <advait.shukla10@gmail.com>
Date: Sat, 22 Aug 2026 17:09:02 -0400
Subject: [PATCH 1/2] Fix sc.pl.paga cax handling with multiple colors

The `cax` parameter was documented as a single axes but indexed as a
sequence (`cax[icolor]`), raising `TypeError: 'Axes' object is not
subscriptable` when plotting multiple colors. Accept either a single
axes (for one colorbar) or a sequence of axes, and raise a clear error
when a single axes is given with multiple colorbars.

Fixes #4318
---
 src/scanpy/plotting/legacy/_tools/paga.py | 15 ++++++++--
 tests/plotting/legacy/test_paga.py        | 35 +++++++++++++++++++++++
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/src/scanpy/plotting/legacy/_tools/paga.py b/src/scanpy/plotting/legacy/_tools/paga.py
index 015eaaa003..7295687c55 100644
--- a/src/scanpy/plotting/legacy/_tools/paga.py
+++ b/src/scanpy/plotting/legacy/_tools/paga.py
@@ -381,7 +381,7 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
     pos: np.ndarray | Path | str | None = None,
     normalize_to_color: bool = False,
     cmap: str | Colormap | None = None,
-    cax: Axes | None = None,
+    cax: Axes | Sequence[Axes] | None = None,
     colorbar=None,  # TODO: this seems to be unused
     cb_kwds: Mapping[str, Any] = frozendict({}),
     frameon: bool | None = None,
@@ -491,7 +491,8 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
     cmap
         The color map.
     cax
-        A matplotlib axes object for a potential colorbar.
+        A matplotlib axes object, or a sequence of axes (one per color), for
+        a potential colorbar.
     cb_kwds
         Keyword arguments for :class:`~matplotlib.colorbar.Colorbar`,
         for instance, `ticks`.
@@ -685,8 +686,16 @@ def is_flat(x):
                     rectangle = [left, bottom, width, height]
                     fig = plt.gcf()
                     ax_cb = fig.add_axes(rectangle)
-                else:
+                elif isinstance(cax, (list, tuple, np.ndarray)):
                     ax_cb = cax[icolor]
+                else:
+                    if sum(colorbars) > 1:
+                        msg = (
+                            "`cax` must be a sequence of axes (one per color) "
+                            "when multiple colorbars are requested."
+                        )
+                        raise ValueError(msg)
+                    ax_cb = cax
 
                 _ = plt.colorbar(
                     sct,
diff --git a/tests/plotting/legacy/test_paga.py b/tests/plotting/legacy/test_paga.py
index 3c1f9e65ff..f1c215121d 100644
--- a/tests/plotting/legacy/test_paga.py
+++ b/tests/plotting/legacy/test_paga.py
@@ -3,9 +3,13 @@
 from functools import partial
 from importlib.util import find_spec
 
+import numpy as np
+import pandas as pd
 import pytest
 from matplotlib import colormaps
+from matplotlib import pyplot as plt
 from packaging.version import Version
+from scipy import sparse
 
 import scanpy as sc
 from scanpy._compat import pkg_version
@@ -96,3 +100,34 @@ def test_paga_compare(plot_cmp):
     sc.pl.paga_compare(pbmc, basis="umap", show=False)
 
     plot_cmp("paga_compare_pbmc3k")
+
+
+def test_paga_cax() -> None:
+    # Tests that https://github.com/scverse/scanpy/issues/4318 is fixed
+    rng = np.random.default_rng(0)
+    adata = sc.AnnData(rng.random((80, 20)))
+    adata.obs["group"] = pd.Categorical(rng.choice(["a", "b", "c", "d", "e"], 80))
+
+    k = 5
+    rows = np.array([0, 1, 1, 2, 2, 3, 3, 4, 4, 0])
+    cols = np.array([1, 0, 2, 1, 3, 2, 4, 3, 0, 4])
+    connectivities = sparse.csr_matrix(  # noqa: TID251
+        (np.ones(len(rows)), (rows, cols)), shape=(k, k)
+    )
+    adata.uns["paga"] = {
+        "groups": "group",
+        "connectivities": connectivities,
+        "connectivities_tree": connectivities.copy(),
+    }
+    pos = rng.random((k, 2))
+
+    # a single `cax` works for a single colorbar
+    _, cax = plt.subplots()
+    sc.pl.paga(adata, color=adata.var_names[0], cax=cax, pos=pos, show=False)
+
+    # a single `cax` with multiple colorbars raises a clear error
+    _, cax = plt.subplots()
+    with pytest.raises(ValueError, match="sequence of axes"):
+        sc.pl.paga(
+            adata, color=adata.var_names[:2].tolist(), cax=cax, pos=pos, show=False
+        )

From 1a1edb2049a8fab576dc7fc17fea161d9f47f6e8 Mon Sep 17 00:00:00 2001
From: Advait <advait.shukla10@gmail.com>
Date: Sat, 22 Aug 2026 17:09:32 -0400
Subject: [PATCH 2/2] Add release note for pl.paga cax fix

---
 docs/release-notes/4319.fix.md | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 docs/release-notes/4319.fix.md

diff --git a/docs/release-notes/4319.fix.md b/docs/release-notes/4319.fix.md
new file mode 100644
index 0000000000..9a540bb61b
--- /dev/null
+++ b/docs/release-notes/4319.fix.md
@@ -0,0 +1 @@
+Fix `sc.pl.paga` `TypeError` when a single `cax` was passed together with multiple `color`s {smaller}`Advait Shukla`