Repository: google/jax

Issue:
`jax.image.resize` failing `linear` interpolation if output shape is `0`
```python
import jax
from jax import numpy as np

image = np.ones((1,))
b = jax.image.resize(image, shape=(0,), method='linear')
```
Gives
```python
---------------------------------------------------------------------------
UnfilteredStackTrace                      Traceback (most recent call last)
<ipython-input-2-2ad1b76cc7df> in <module>()
      4 image = np.ones((1,))
----> 5 b = jax.image.resize(image, shape=(0,), method='linear')

20 frames
UnfilteredStackTrace: ZeroDivisionError: float division by zero

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:

ZeroDivisionError                         Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/jax/_src/image/scale.py in compute_weight_mat(input_size, output_size, scale, translation, kernel, antialias)
     48                        kernel: Callable,
     49                        antialias: bool):
---> 50   inv_scale = 1. / scale
     51   # When downsampling the kernel should be scaled since we want to low pass
     52   # filter and interpolate, but when upsampling it should not be since we only

ZeroDivisionError: float division by zero
```
Although arguably the output should just be an empty array of the same type as input (and it works with `method="nearest"`).


Discussion:

Resolving patch:
diff --git a/jax/_src/image/scale.py b/jax/_src/image/scale.py
--- a/jax/_src/image/scale.py
+++ b/jax/_src/image/scale.py
@@ -253,7 +253,8 @@ def _resize(image, shape: Sequence[int], method: Union[str, ResizeMethod],
   # since all of the current resize methods (kernels) are interpolating, so the
   # output = input under an identity warp.
   spatial_dims = tuple(np.nonzero(np.not_equal(image.shape, shape))[0])
-  scale = [float(shape[d]) / image.shape[d] for d in spatial_dims]
+  scale = [1.0 if shape[d] == 0 else float(shape[d]) / image.shape[d]
+           for d in spatial_dims]
   return _scale_and_translate(image, shape, spatial_dims,
                               scale, [0.] * len(spatial_dims), kernel,
                               antialias, precision)


Test patch:
diff --git a/tests/image_test.py b/tests/image_test.py
--- a/tests/image_test.py
+++ b/tests/image_test.py
@@ -178,6 +178,29 @@ def testResizeGradients(self, dtype, image_shape, target_shape, method,
                      antialias=antialias)
     jtu.check_grads(jax_fn, args_maker(), order=2, rtol=1e-2, eps=1.)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+       {"testcase_name": "_shape={}_target={}_method={}_antialias={}".format(
+          jtu.format_shape_dtype_string(image_shape, dtype),
+          jtu.format_shape_dtype_string(target_shape, dtype), method,
+          antialias),
+        "dtype": dtype, "image_shape": image_shape,
+        "target_shape": target_shape,
+        "method": method, "antialias": antialias}
+       for dtype in [np.float32]
+       for image_shape, target_shape in [
+         ([1], [0]),
+         ([5, 5], [5, 0]),
+         ([5, 5], [0, 1]),
+         ([5, 5], [0, 0])
+       ]
+       for method in ["nearest", "linear", "lanczos3", "lanczos5", "cubic"]
+       for antialias in [False, True]))
+  def testResizeEmpty(self, dtype, image_shape, target_shape, method, antialias):
+    # Regression test for https://github.com/google/jax/issues/7586
+    image = np.ones(image_shape, dtype)
+    out = jax.image.resize(image, shape=target_shape, method=method, antialias=antialias)
+    self.assertArraysEqual(out, jnp.zeros(target_shape, dtype))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_target={}_method={}".format(
          jtu.format_shape_dtype_string(image_shape, dtype),