Repository: google/jax

Issue:
`lax.conv_general_dilated` will crash when `padding` has some negative
`lax.conv_general_dilated` will crash when `padding` has some negative

```py
import jax
def fn(lhs, rhs):
    window_strides = [1, 1]
    padding = [[5, 16], [-5, 1]]
    return jax.lax.conv_general_dilated(lhs, rhs, window_strides, padding)

mykey = jax.random.PRNGKey(83725540)

lhs = jax.numpy.array([[[[-7.0950737, 20.612083, 53.864067], [-12.116554, -16.742317, -7.050851 ]]]])
rhs = jax.numpy.array([[[[29.074692, 40.518192], [21.685156, 41.867775]]]])

fn(lhs, rhs)
```
```
F external/org_tensorflow/tensorflow/compiler/xla/window_util.cc:250] Check failed: bound >= 0 (0 vs. -1)
[1]   abort (core dumped)
```

Besides, `jax.lax.conv_general_dilated_patches` and `jax.lax.conv_with_general_padding` also have such issue

```py
import jax
def fn(lhs):
    filter_shape = [1, 1]
    window_strides = [1, 1]
    padding = [[-4, -2], [-3, -4]]
    return jax.lax.conv_general_dilated_patches(lhs=lhs, filter_shape=filter_shape, window_strides=window_strides, padding=padding)

mykey = jax.random.PRNGKey(19446744)
lhs = jax.random.uniform(mykey, [1, 2, 3, 4], jax.numpy.float32)
fn(lhs)
```

```
F external/org_tensorflow/tensorflow/compiler/xla/window_util.cc:250] Check failed: bound >= 0 (0 vs. -3)
[1]   abort (core dumped)
```

```py
import jax
def fn(lhs, rhs):
    window_strides = [1, 1]
    padding = [[-5, 0], [-16, -2]]
    lhs_dilation = [1, 1]
    rhs_dilation = [1, 1]
    return jax.lax.conv_with_general_padding(lhs, rhs, window_strides=window_strides, padding=padding, lhs_dilation=lhs_dilation, rhs_dilation=rhs_dilation)

mykey = jax.random.PRNGKey(30297051)
lhs = jax.random.uniform(mykey, [3, 2, 3, 1], jax.numpy.float32)
mykey = jax.random.PRNGKey(34585797)
rhs = jax.random.uniform(mykey, [2, 2, 1, 2], jax.numpy.float32)

fn(lhs, rhs)
```

```
F external/org_tensorflow/tensorflow/compiler/xla/window_util.cc:250] Check failed: bound >= 0 (0 vs. -2)
[1]   abort (core dumped)
```


Discussion:
Hi @zhangqiaorjc, is this [comment](https://github.com/google/jax/issues/11521#issuecomment-1188254569) in a wrong place?

If so, I agree that a check in python is much better than core dumped in cpp!
https://github.com/google/jax/pull/11544 should fix the issue
Actually, it looks like negative padding is valid after all; from https://www.tensorflow.org/xla/operation_semantics#convwithgeneralpadding_convolution:

> The padding argument specifies the amount of zero padding to be applied to the base area. The amount of padding can be negative -- the absolute value of negative padding indicates the number of elements to remove from the specified dimension before doing the convolution.

I think the issue in your example is that the negative value is too large for the size of the specified axis. This is going to be a little more involved to check correctly on the Python side...

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
@@ -812,6 +812,9 @@ def conv_shape_tuple(lhs_shape, rhs_shape, strides, pads, batch_group_count=1):
 
   lhs_padded = np.add(lhs_shape[2:], np.sum(np.array(pads).reshape(-1, 2),
                                               axis=1))
+  if np.any(lhs_padded < 0):
+    raise ValueError("Negative padding is larger than the size of the corresponding dimension: "
+                     f"got padding={pads} for lhs_shape[2:]={lhs_shape[2:]}")
   out_space = core.stride_shape(lhs_padded, rhs_shape[2:], strides)
   out_space = np.maximum(0, out_space)
   if batch_group_count > 1:


Test patch: