Repository: google/jax

Issue:
`lax.pad` reverse-mode AD fails on scalar inputs
https://colab.research.google.com/gist/romanngg/8284adda14131eca8d63d69929044277/lax_pad_reverse_mode_fail.ipynb

```python
def f(x): 
  return lax.pad(x, 0., ())

jacrev(f)(1.)
```
gives
```python
---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation         Traceback (most recent call last)
[/usr/lib/python3.7/runpy.py](https://localhost:8080/#) in _run_module_as_main(***failed resolving arguments***)
    192     return _run_code(code, main_globals, None,
--> 193                      "__main__", mod_spec)
    194 

36 frames
JaxStackTraceBeforeTransformation: ValueError: not enough values to unpack (expected 3, got 0)

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

UnfilteredStackTrace                      Traceback (most recent call last)
UnfilteredStackTrace: ValueError: not enough values to unpack (expected 3, got 0)

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/jax/_src/api.py](https://localhost:8080/#) in jacfun(*args, **kwargs)
   1170       y, pullback, aux = _vjp(f_partial, *dyn_args, has_aux=True)
   1171     tree_map(partial(_check_output_dtype_jacrev, holomorphic), y)
-> 1172     jac = vmap(pullback)(_std_basis(y))
   1173     jac = jac[0] if isinstance(argnums, int) else jac
   1174     example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args

ValueError: not enough values to unpack (expected 3, got 0)
```


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2893,7 +2893,7 @@ def _pad_transpose(t, operand, padding_value, *, padding_config):
     t_operand = ad_util.Zero(operand.aval) if ad.is_undefined_primal(operand) else None
     t_padv = ad_util.Zero(padding_value.aval) if ad.is_undefined_primal(padding_value) else None
   else:
-    lo, hi, interior = zip(*padding_config)
+    lo, hi, interior = util.unzip3(padding_config)
     total = lambda x: _reduce_sum(x, list(range(t.ndim)))
 
     def t_op():
diff --git a/jax/_src/lax/windowed_reductions.py b/jax/_src/lax/windowed_reductions.py
--- a/jax/_src/lax/windowed_reductions.py
+++ b/jax/_src/lax/windowed_reductions.py
@@ -418,7 +418,7 @@ def reduce_window_shape_tuple(operand_shape, window_dimensions, window_strides,
     operand_shape = lax._dilate_shape(operand_shape, base_dilation)
   if window_dilation is not None:
     window_dimensions = lax._dilate_shape(window_dimensions, window_dilation)
-  pads_lo, pads_hi = [(), ()] if len(padding) == 0 else zip(*padding)
+  pads_lo, pads_hi = util.unzip2(padding)
   operand_padded = core.sum_shapes(operand_shape, pads_lo, pads_hi)
   return core.stride_shape(operand_padded, window_dimensions, window_strides)
 


Test patch:
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -502,9 +502,12 @@ def testReshapeGrad(self, arg_shape, out_shape, permutation, dtype):
       {"testcase_name": "_inshape={}_pads={}"
        .format(jtu.format_shape_dtype_string(shape, dtype), pads),
        "shape": shape, "dtype": dtype, "pads": pads}
-      for shape in [(2, 3)]
       for dtype in float_dtypes
-      for pads in [[(1, 2, 1), (0, 1, 0)], [(-1, 0, 0), (-1, 0, 2)]]))
+      for shape, paddings in [
+        [(), [()]],
+        ((2, 3), [[(1, 2, 1), (0, 1, 0)], [(-1, 0, 0), (-1, 0, 2)]]),
+      ]
+      for pads in paddings))
   def testPadGrad(self, shape, dtype, pads):
     rng = jtu.rand_small(self.rng())
     operand = rng(shape, dtype)