Repository: google/jax
Issue: Low-level error in forward-mode AD of `lax.reduce_window_sum_p`

https://colab.research.google.com/gist/icml2022anon/bdc92ea67396a90eef22eaf80ab2f9f3/reduce_window_sum_fwd_bug.ipynb

```python
f = lambda x: jax.lax.reduce_window_sum_p.bind(x, **{
            'base_dilation': (1, 1, 1),
            'padding': ((0, 0), (0, 0), (0, 0)),
            'window_dilation': (3, 1, 1),
            'window_dimensions': (3, 1, 1),
            'window_strides': (1, 2, 1)
        })
```
works on forward pass and reverse mode, but `jax.jacfwd(f)(np.ones((5, 3, 2)))` causes
```python
---------------------------------------------------------------------------
UnfilteredStackTrace                      Traceback (most recent call last)
[<ipython-input-6-f0f9409ecd4d>](https://localhost:8080/#) in <module>()
----> 1 jax.jacfwd(f)(np.ones((5, 3, 2)))

27 frames
UnfilteredStackTrace: RuntimeError: UNKNOWN: <unknown>:0: error: Expected array argument for transpose, but got primitive_type_invalid[].: 
<unknown>:0: note: see current operation: "func.return"(%0) : (tensor<?x2x2x30xf32>) -> ()

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:

RuntimeError                              Traceback (most recent call last)
[/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py](https://localhost:8080/#) in transpose(a, axes)
    494   _check_arraylike("transpose", a)
    495   axes = np.arange(ndim(a))[::-1] if axes is None else axes
--> 496   return lax.transpose(a, axes)
    497 
    498 

RuntimeError: UNKNOWN: <unknown>:0: error: Expected array argument for transpose, but got primitive_type_invalid[].: 
<unknown>:0: note: see current operation: "func.return"(%0) : (tensor<?x2x2x30xf32>) -> ()
```


Changed files:
jax/experimental/jax2tf/shape_poly.py, tests/lax_test.py, jax/core.py

diff --git a/jax/core.py b/jax/core.py
index 7ede76b0b..2170cf972 100644
--- a/jax/core.py
+++ b/jax/core.py
@@ -1511,7 +1511,11 @@ class DimensionHandler:
     return sz1 // sz2
 
   def stride(self, d: DimSize, window_size: DimSize, window_stride: DimSize) -> DimSize:
-    """(d - window_size) // window_stride + 1"""
+    """(d - window_size) // window_stride + 1.
+
+    If d == 0 or window_size > d, returns 0.
+    """
+    if d == 0 or window_size > d: return 0
     return (d - window_size) // window_stride + 1
 
   def dilate(self, d: DimSize, dilation: int) -> DimSize:
diff --git a/jax/experimental/jax2tf/shape_poly.py b/jax/experimental/jax2tf/shape_poly.py
index 1fee18d18..4465fea2b 100644
--- a/jax/experimental/jax2tf/shape_poly.py
+++ b/jax/experimental/jax2tf/shape_poly.py
@@ -442,6 +442,7 @@ class DimensionHandlerPoly(core.DimensionHandler):
   def stride(self, d: DimSize, window_size: DimSize, window_stride: DimSize) -> DimSize:
     """Implements `(d - window_size) // window_stride + 1`"""
     try:
+      # TODO(necula): check for d == 0 or window_size > d and return 0.
       q, r = _ensure_poly(d - window_size).divmod(window_stride)
       return q + 1
     except InconclusiveDimensionOperation as e:
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 7502eb47c..df3cd3508 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1931,6 +1931,17 @@ class LaxTest(jtu.JaxTestCase):
     # to https://www.tensorflow.org/xla/operation_semantics#reducewindow.
     self.assertEqual(shape, result.shape)
 
+  def testReduceWindowWithEmptyOutput(self):
+    # https://github.com/google/jax/issues/10315
+    shape = (5, 3, 2)
+    operand, padding, strides = np.ones(shape), 'VALID', (1,) * len(shape)
+    out = jax.eval_shape(lambda x: lax.reduce_window(x, 0., lax.add, padding=padding,
+                         window_strides=strides,
+                         window_dimensions=(3, 1, 1),
+                         window_dilation=(3, 1, 1)), operand)
+    self.assertEqual((0, 3, 2), out.shape)
+
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_op={}_shape={}_axis={}_reverse={}"
        .format(op.__name__, jtu.format_shape_dtype_string(shape, dtype), axis,