Repository: google/jax

Issue:
Suspected bug in jax.lax.conv_general_dilated_patches
Trying to implement the solution [here](https://github.com/google/jax/discussions/5968) but ran into the following bug:

```import jax
import jax.numpy as jnp

dev = jax.local_devices()[0]

ftmap_fn = jax.device_put(jnp.ones((1, 384, 384, 256), dtype=jnp.bfloat16), device=dev)
W = 5
stride, padding = (4, 4), (W // 2, W // 2)
jax.lax.conv_general_dilated_patches(ftmap_fn, (5, 5), stride, padding).shape
```
but i found the following error:
```
...
File ~/.local/lib/python3.8/site-packages/jax/_src/dispatch.py:609, in XlaCompiledComputation.from_xla_computation(name, xla_computation, nreps, device, backend, tuple_args, in_avals, out_avals, kept_var_idx)
    606 options.parameter_is_tupled_arguments = tuple_args
    607 with log_elapsed_time(f"Finished XLA compilation of {name} "
    608                       "in {elapsed_time} sec"):
--> 609   compiled = compile_or_get_cached(backend, xla_computation, options)
    610 buffer_counts = (None if len(out_avals) == 1 else
    611                  [aval_to_num_buffers(aval) for aval in out_avals])
    612 execute = _execute_compiled if nreps == 1 else _execute_replicated

File ~/.local/lib/python3.8/site-packages/jax/_src/dispatch.py:578, in compile_or_get_cached(backend, computation, compile_options)
    575   ir_str = (computation if isinstance(computation, str)
    576             else computation.as_hlo_text())
    577   _dump_ir_to_file(module_name, ir_str)
--> 578 return backend_compile(backend, computation, compile_options)

File ~/.local/lib/python3.8/site-packages/jax/_src/profiler.py:206, in annotate_function.<locals>.wrapper(*args, **kwargs)
    203 @wraps(func)
    204 def wrapper(*args, **kwargs):
    205   with TraceAnnotation(name, **decorator_kwargs):
--> 206     return func(*args, **kwargs)
    207   return wrapper

File ~/.local/lib/python3.8/site-packages/jax/_src/dispatch.py:532, in backend_compile(backend, built_c, options)
    528 @profiler.annotate_function
    529 def backend_compile(backend, built_c, options):
    530   # we use a separate function call to ensure that XLA compilation appears
    531   # separately in Python profiling results
--> 532   return backend.compile(built_c, compile_options=options)

RuntimeError: UNKNOWN: -:4:130: error: expected '['
``` 

The error looks very cryptic cannot figure out what it means. Is it an XLA / Jax bug? Thanks!

Environment: TPU-VM v3-8 with software version tf.2.8.0
```
jax==0.3.5
jaxlib==0.3.5
libtpu-nightly==0.1.dev20220407
cloud-tpu-client==0.10
```


Discussion:
The issue is that `padding` should be a list of pairs, one for each dimension. e.g., this works:

```
stride, padding = (4, 4), [(W // 2, W // 2), (W // 2, W //2)]
```

We are apparently missing some input validation that would have caught this sooner. I'll fix that!

Resolving patch:
diff --git a/jax/_src/lax/convolution.py b/jax/_src/lax/convolution.py
--- a/jax/_src/lax/convolution.py
+++ b/jax/_src/lax/convolution.py
@@ -14,6 +14,7 @@
 
 import builtins
 from functools import partial
+import operator
 from typing import Any, List, NamedTuple, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -141,6 +142,15 @@ def conv_general_dilated(
     padding = lax.padtype_to_pads(
         np.take(lhs.shape, lhs_perm)[2:], effective_rhs_shape,  # type: ignore[index]
         window_strides, padding)
+  else:
+    try:
+      padding = tuple((operator.index(lo), operator.index(hi))
+                      for lo, hi in padding)
+    except (ValueError, TypeError) as e:
+      raise ValueError(
+        "padding argument to conv_general_dilated should be a string or a "
+        f"sequence of (low, high) pairs, got {padding}") from e
+
   preferred_element_type = (
       None if preferred_element_type is None else
       dtypes.canonicalize_dtype(np.dtype(preferred_element_type)))


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1060,6 +1060,13 @@ def testConvTransposePaddingList(self):
     c = lax.conv_general_dilated(a[None, None], b[None, None], (1,1), [(0,0),(0,0)], (1,1))
     self.assertAllClose(c, 9 * jnp.ones((1, 1, 26, 26)))
 
+  def testConvInvalidPadding(self):
+    x = jnp.ones((1, 10, 10, 5), dtype=jnp.bfloat16)
+    with self.assertRaisesRegex(ValueError,
+                                r"padding argument.*, got \(3, 3\)"):
+      jax.lax.conv_general_dilated_patches(x, (5, 5), window_strides=(1, 1),
+                                           padding=(3, 3))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_lhs_shape={}_rhs_shape={}_precision={}".format(
           jtu.format_shape_dtype_string(lhs_shape, dtype),