Repository: dask/dask

Issue:
dask.array.pad does not chunk up padded region
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:

**Minimal Complete Verifiable Example**:

```python
import dask.array

array = dask.array.ones( (61440, 30757), chunks=(4000, 4000))
dask.array.pad(array, ((0, 88970 - 61440), (0,0)), mode="constant", constant_values=0)
```

<img width="449" alt="Image" src="https://github.com/user-attachments/assets/61b8e385-c58d-4068-ad0e-bf04fdd7e469" />


Discussion:

Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index 9265bc1b6..6170691ee 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -40,9 +40,11 @@ def _unpack_collections(o):
 
 
 class Expr:
-    _parameters = []  # type: ignore
+    _parameters: list[str] = []
     _defaults: dict[str, Any] = {}
-    _instances = weakref.WeakValueDictionary()  # type: ignore
+    _instances: weakref.WeakValueDictionary[str, Expr] = weakref.WeakValueDictionary()
+
+    operands: list
 
     def __new__(cls, *args, **kwargs):
         operands = list(args)
@@ -99,16 +101,21 @@ class Expr:
         return os.linesep.join(self._tree_repr_lines())
 
     def analyze(self, filename: str | None = None, format: str | None = None) -> None:
+        from dask.dataframe.dask_expr._expr import Expr as DFExpr
         from dask.dataframe.dask_expr.diagnostics import analyze
 
-        return analyze(self, filename=filename, format=format)  # type: ignore
+        if not isinstance(self, DFExpr):
+            raise TypeError(
+                "analyze is only supported for dask.dataframe.Expr objects."
+            )
+        return analyze(self, filename=filename, format=format)
 
     def explain(
         self, stage: OptimizerStage = "fused", format: str | None = None
     ) -> None:
         from dask.dataframe.dask_expr.diagnostics import explain
 
-        return explain(self, stage, format)  # type: ignore
+        return explain(self, stage, format)
 
     def pprint(self):
         for line in self._tree_repr_lines():
@@ -266,7 +273,7 @@ class Expr:
             # Rewrite all of the children
             new_operands = []
             changed = False
-            for operand in expr.operands:  # type: ignore
+            for operand in expr.operands:
                 if isinstance(operand, Expr):
                     new = operand.rewrite(kind=kind, rewritten=rewritten)
                     rewritten[operand._name] = new
@@ -333,7 +340,7 @@ class Expr:
             # Rewrite all of the children
             new_operands = []
             changed = False
-            for operand in expr.operands:  # type: ignore
+            for operand in expr.operands:
                 if isinstance(operand, Expr):
                     # Bandaid for now, waiting for Singleton
                     dependents[operand._name].append(weakref.ref(expr))
@@ -354,6 +361,14 @@ class Expr:
 
         return expr
 
+    def optimize(self, fuse: bool = False) -> Expr:
+        stage: OptimizerStage = "fused" if fuse else "simplified-physical"
+
+        return optimize_until(self, stage)
+
+    def fuse(self) -> Expr:
+        return self
+
     def simplify(self) -> Expr:
         expr = self
         seen = set()
@@ -396,7 +411,7 @@ class Expr:
         # Lower all children
         new_operands = []
         changed = False
-        for operand in out.operands:  # type: ignore
+        for operand in out.operands:
             if isinstance(operand, Expr):
                 new = operand.lower_once(lowered)
                 if new._name != operand._name:
@@ -442,11 +457,11 @@ class Expr:
         return
 
     @functools.cached_property
-    def _funcname(self):
+    def _funcname(self) -> str:
         return funcname(type(self)).lower()
 
     @functools.cached_property
-    def _name(self):
+    def _name(self) -> str:
         return self._funcname + "-" + _tokenize_deterministic(*self.operands)
 
     @property
@@ -563,7 +578,7 @@ class Expr:
 
         changed = False
         new_operands = []
-        for i, operand in enumerate(self.operands):  # type: ignore
+        for i, operand in enumerate(self.operands):
             if i < len(self._parameters) and self._parameters[i] in substitutions:
                 new_operands.append(substitutions[self._parameters[i]])
                 changed = True
@@ -733,3 +748,61 @@ def collect_dependents(expr) -> defaultdict:
             stack.append(dep)
             dependents[dep._name].append(weakref.ref(node))
     return dependents
+
+
+def optimize(expr: Expr, fuse: bool = True) -> Expr:
+    """High level query optimization
+
+    This leverages three optimization passes:
+
+    1.  Class based simplification using the ``_simplify`` function and methods
+    2.  Blockwise fusion
+
+    Parameters
+    ----------
+    expr:
+        Input expression to optimize
+    fuse:
+        whether or not to turn on blockwise fusion
+
+    See Also
+    --------
+    simplify
+    optimize_blockwise_fusion
+    """
+    stage: OptimizerStage = "fused" if fuse else "simplified-physical"
+
+    return optimize_until(expr, stage)
+
+
+def optimize_until(expr: Expr, stage: OptimizerStage) -> Expr:
+    result = expr
+    if stage == "logical":
+        return result
+
+    # Simplify
+    expr = result.simplify()
+    if stage == "simplified-logical":
+        return expr
+
+    # Manipulate Expression to make it more efficient
+    expr = expr.rewrite(kind="tune", rewritten={})
+    if stage == "tuned-logical":
+        return expr
+
+    # Lower
+    expr = expr.lower_completely()
+    if stage == "physical":
+        return expr
+
+    # Simplify again
+    expr = expr.simplify()
+    if stage == "simplified-physical":
+        return expr
+
+    # Final graph-specific optimizations
+    expr = expr.fuse()
+    if stage == "fused":
+        return expr
+
+    raise ValueError(f"Stage {stage!r} not supported.")
diff --git a/dask/array/creation.py b/dask/array/creation.py
index 0f69d165c..7a1edb2a0 100644
--- a/dask/array/creation.py
+++ b/dask/array/creation.py
@@ -968,7 +968,7 @@ def expand_pad_value(array, pad_value):
     return pad_value
 
 
-def get_pad_shapes_chunks(array, pad_width, axes):
+def get_pad_shapes_chunks(array, pad_width, axes, mode):
     """
     Helper function for finding shapes and chunks of end pads.
     """
@@ -979,7 +979,12 @@ def get_pad_shapes_chunks(array, pad_width, axes):
     for d in axes:
         for i in range(2):
             pad_shapes[i][d] = pad_width[d][i]
-            pad_chunks[i][d] = (pad_width[d][i],)
+            if mode != "constant" or pad_width[d][i] == 0:
+                pad_chunks[i][d] = (pad_width[d][i],)
+            else:
+                pad_chunks[i][d] = normalize_chunks(
+                    (max(pad_chunks[i][d]),), (pad_width[d][i],)
+                )[0]
 
     pad_shapes = [tuple(s) for s in pad_shapes]
     pad_chunks = [tuple(c) for c in pad_chunks]
@@ -1021,7 +1026,9 @@ def pad_edge(array, pad_width, mode, **kwargs):
 
     result = array
     for d in range(array.ndim):
-        pad_shapes, pad_chunks = get_pad_shapes_chunks(result, pad_width, (d,))
+        pad_shapes, pad_chunks = get_pad_shapes_chunks(
+            result, pad_width, (d,), mode=mode
+        )
         pad_arrays = [result, result]
 
         if mode == "constant":
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index fe7390bdb..032717ba5 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -61,6 +61,8 @@ from dask.utils import (
     random_state_data,
 )
 
+optimize = core.optimize
+
 
 class Expr(core.Expr):
     """Primary class for all Expressions
@@ -522,6 +524,9 @@ class Expr(core.Expr):
             predicate = parent.predicate.substitute(self, self.frame)
         return type(self)(self.frame[predicate], *self.operands[1:])
 
+    def fuse(self):
+        return optimize_blockwise_fusion(self)
+
 
 class Literal(Expr):
     """Represent a literal (known) value as an `Expr`"""
@@ -3074,7 +3079,7 @@ class _DelayedExpr(Expr):
 
 
 class DelayedsExpr(Expr):
-    _parameters = []  # type: ignore
+    _parameters = []
 
     def __init__(self, *delayed_objects):
         self.operands = delayed_objects
@@ -3108,64 +3113,6 @@ def normalize_expression(expr):
     return expr._name
 
 
-def optimize_until(expr: Expr, stage: core.OptimizerStage) -> Expr:
-    result = expr
-    if stage == "logical":
-        return result
-
-    # Simplify
-    expr = result.simplify()  # type: ignore
-    if stage == "simplified-logical":
-        return expr
-
-    # Manipulate Expression to make it more efficient
-    expr = expr.rewrite(kind="tune", rewritten={})
-    if stage == "tuned-logical":
-        return expr
-
-    # Lower
-    expr = expr.lower_completely()  # type: ignore
-    if stage == "physical":
-        return expr
-
-    # Simplify again
-    expr = expr.simplify()  # type: ignore
-    if stage == "simplified-physical":
-        return expr
-
-    # Final graph-specific optimizations
-    expr = optimize_blockwise_fusion(expr)
-    if stage == "fused":
-        return expr
-
-    raise ValueError(f"Stage {stage!r} not supported.")
-
-
-def optimize(expr: Expr, fuse: bool = True) -> Expr:
-    """High level query optimization
-
-    This leverages three optimization passes:
-
-    1.  Class based simplification using the ``_simplify`` function and methods
-    2.  Blockwise fusion
-
-    Parameters
-    ----------
-    expr:
-        Input expression to optimize
-    fuse:
-        whether or not to turn on blockwise fusion
-
-    See Also
-    --------
-    simplify
-    optimize_blockwise_fusion
-    """
-    stage: core.OptimizerStage = "fused" if fuse else "simplified-physical"
-
-    return optimize_until(expr, stage)
-
-
 def is_broadcastable(dfs, s):
     """
     This Series is broadcastable against another dataframe in the sequence
diff --git a/dask/dataframe/dask_expr/diagnostics/_explain.py b/dask/dataframe/dask_expr/diagnostics/_explain.py
index 1d5beebc3..5912b82fe 100644
--- a/dask/dataframe/dask_expr/diagnostics/_explain.py
+++ b/dask/dataframe/dask_expr/diagnostics/_explain.py
@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from dask._expr import OptimizerStage
-from dask.dataframe.dask_expr._expr import Expr, Projection, optimize_until
+from dask._expr import Expr as BaseExpr
+from dask._expr import OptimizerStage, optimize_until
+from dask.dataframe.dask_expr._expr import Expr, Projection
 from dask.dataframe.dask_expr._merge import Merge
 from dask.dataframe.dask_expr.io.parquet import ReadParquet
 from dask.utils import funcname, import_required
@@ -17,7 +18,7 @@ STAGE_LABELS: dict[OptimizerStage, str] = {
 
 
 def explain(
-    expr: Expr, stage: OptimizerStage = "fused", format: str | None = None
+    expr: BaseExpr, stage: OptimizerStage = "fused", format: str | None = None
 ) -> None:
     graphviz = import_required(
         "graphviz", "graphviz is a required dependency for using the explain method."
@@ -34,7 +35,7 @@ def explain(
     expr = optimize_until(expr, stage)
 
     seen = set(expr._name)
-    stack = [expr]
+    stack: list[BaseExpr] = [expr]
 
     while stack:
         node = stack.pop()
@@ -73,7 +74,7 @@ def _add_graphviz_edges(explain_info, graph):
         graph.edge(dep, name)
 
 
-def _explain_info(expr: Expr):
+def _explain_info(expr: BaseExpr):
     return {
         "name": expr._name,
         "label": funcname(type(expr)),
@@ -82,7 +83,10 @@ def _explain_info(expr: Expr):
     }
 
 
-def _explain_details(expr: Expr):
+def _explain_details(expr: BaseExpr):
+
+    if not isinstance(expr, Expr):
+        return {}
     details = {"npartitions": expr.npartitions}
 
     if isinstance(expr, Merge):
@@ -96,10 +100,10 @@ def _explain_details(expr: Expr):
     return details
 
 
-def _explain_dependencies(expr: Expr) -> list[tuple[str, str]]:
+def _explain_dependencies(expr: BaseExpr) -> list[tuple[str, str]]:
     dependencies = []
     for i, operand in enumerate(expr.operands):
-        if not isinstance(operand, Expr):
+        if not isinstance(operand, BaseExpr):
             continue
         param = expr._parameters[i] if i < len(expr._parameters) else ""
         dependencies.append((str(param), operand._name))


Test patch:
diff --git a/dask/array/tests/test_creation.py b/dask/array/tests/test_creation.py
index fef3272e8..9887886c9 100644
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -991,6 +991,23 @@ def test_pad_udf(kwargs):
     assert_eq(np_r, da_r)
 
 
+def test_pad_constant_chunksizes():
+    array = dask.array.ones((10, 10), chunks=(1, 1))
+    result = dask.array.pad(
+        array, ((0, 16 - 10), (0, 0)), mode="constant", constant_values=0
+    )
+    assert tuple(map(max, result.chunks)) == (1, 1)
+    assert_eq(
+        result,
+        np.pad(
+            array.compute(),
+            mode="constant",
+            constant_values=0,
+            pad_width=((0, 16 - 10), (0, 0)),
+        ),
+    )
+
+
 def test_auto_chunks():
     with dask.config.set({"array.chunk-size": "50 MiB"}):
         x = da.ones((10000, 10000))