Merge pull request jax-ml#24896 from hawkinsp:pad

PiperOrigin-RevId: 696568236

Author: Google-ML-Automation

jax/_src/numpy/lax_numpy.py

@@ -4048,15 +4048,37 @@ def _check_no_padding(axis_padding: tuple[Any, Any], mode: str):
 
 def _pad_constant(array: Array, pad_width: PadValue[int], constant_values: Array) -> Array:
   nd = ndim(array)
-  constant_values = broadcast_to(constant_values, (nd, 2))
   constant_values = lax_internal._convert_element_type(
       constant_values, array.dtype, dtypes.is_weakly_typed(array))
+  constant_values_nd = ndim(constant_values)
+
+  if constant_values_nd == 0:
+    widths = [(low, high, 0) for (low, high) in pad_width]
+    return lax.pad(array, constant_values, widths)
+
+  if constant_values_nd == 1:
+    if constant_values.shape[-1] == 1:
+      widths = [(low, high, 0) for (low, high) in pad_width]
+      return lax.pad(array, squeeze(constant_values), widths)
+    elif constant_values.shape[-1] == 2:
+      widths = [(low, 0, 0) for (low, _) in pad_width]
+      array = lax.pad(array, constant_values[0], widths)
+      widths = [(0, high, 0) for (_, high) in pad_width]
+      return lax.pad(array, constant_values[1], widths)
+    else:
+      raise ValueError("jnp.pad: constant_values has unsupported shape "
+                      f"{constant_values.shape}. If the shape is 1D or 2D, the "
+                      "last dimension must be of size 1 or 2.")
+
+  constant_values = broadcast_to(constant_values, (nd, 2))
   for i in range(nd):
     widths = [(0, 0, 0)] * nd
-    widths[i] = (pad_width[i][0], 0, 0)
-    array = lax.pad(array, constant_values[i, 0], widths)
-    widths[i] = (0, pad_width[i][1], 0)
-    array = lax.pad(array, constant_values[i, 1], widths)
+    if pad_width[i][0] != 0:
+      widths[i] = (pad_width[i][0], 0, 0)
+      array = lax.pad(array, constant_values[i, 0], widths)
+    if pad_width[i][1] != 0:
+      widths[i] = (0, pad_width[i][1], 0)
+      array = lax.pad(array, constant_values[i, 1], widths)
   return array