Repository: google/jax
Issue: nans in gradient of scale_and_translate

I'd like to take gradients of `scale_and_translate` wrt the translation parameters.  Currently the gradients are `nan` whenever the translation is >1.  Example: 

```python
import jax
import jax.numpy as jnp
from jax.image import scale_and_translate

xin = jnp.r_[:16].reshape(4, 4) + 1.
xin = jax.device_put(xin)

def f(x):
    out = scale_and_translate(xin, shape=xin.shape, spatial_dims=(0, 1),
                              scale=jnp.array([1.0, 1.0]),
                              translation=jnp.array([x[0], x[1]]),
                              method='linear', antialias=True)
    return out

f((1.1, 1.2))  # raises an error if JAX_DEBUG_NANS=1

g = lambda x: jnp.sum(f(x))

print(jax.grad(g)( (1.1, 1.2)))  # nans
```

If `JAX_DEBUG_NANS=1` is set, the script throws an exception at the evaluation of `f`:
```
Traceback (most recent call last):
  File "resize.py", line 15, in <module>
    f((1.1, 1.2))  # raises an error if JAX_DEBUG_NANS=1
  File "resize.py", line 9, in f
    out = scale_and_translate(xin, shape=xin.shape, spatial_dims=(0, 1),
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/_src/image/scale.py", line 219, in scale_and_translate
    return _scale_and_translate(image, shape, spatial_dims, scale, translation,
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/_src/image/scale.py", line 92, in _scale_and_translate
    w = compute_weight_mat(m, n, scale[i], translation[i],
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/_src/image/scale.py", line 69, in compute_weight_mat
    weights / total_weight_sum, 0)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 5666, in deferring_binary_op
    return binary_op(self, other)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 569, in true_divide
    return lax.div(x1, x2)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/_src/lax/lax.py", line 352, in div
    return div_p.bind(x, y)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/core.py", line 264, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/core.py", line 606, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/interpreters/xla.py", line 232, in apply_primitive
    return compiled_fun(*args)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/interpreters/xla.py", line 351, in _execute_compiled_primitive
    check_special(prim.name, out_bufs)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/interpreters/xla.py", line 370, in check_special
    _check_special(name, buf.xla_shape(), buf)
  File "/home/pfister/miniconda3/envs/scico/lib/python3.8/site-packages/jax/interpreters/xla.py", line 376, in _check_special
    raise FloatingPointError(f"invalid value (nan) encountered in {name}")
FloatingPointError: invalid value (nan) encountered in div
```

The issue is in this block: 
https://github.com/google/jax/blob/c85e835951c68b4383dbc9ef25bef99ea996fd84/jax/_src/image/scale.py#L66-L69
in particular, the computaiton `weights/total_weight_sum` can involve a divide by 0, even though those locations get set to zero in the course of the `jnp.where`.

I'll submit a PR shortly with the "where inside of where" fix.


Changed files:
tests/image_test.py, jax/_src/image/scale.py

diff --git a/jax/_src/image/scale.py b/jax/_src/image/scale.py
index f019107ed..b63fd98fc 100644
--- a/jax/_src/image/scale.py
+++ b/jax/_src/image/scale.py
@@ -24,9 +24,8 @@ import numpy as np
 
 def _fill_lanczos_kernel(radius, x):
   y = radius * jnp.sin(np.pi * x) * jnp.sin(np.pi * x / radius)
-  with np.errstate(divide='ignore', invalid='ignore'):
-    out = y / (np.pi ** 2 * x ** 2)
-  out = jnp.where(x <= 1e-3, 1., out)
+  #  out = y / (np.pi ** 2 * x ** 2) where x >1e-3, 1 otherwise
+  out = jnp.where(x > 1e-3, jnp.divide(y, jnp.where(x != 0, np.pi**2 * x**2, 1)), 1)
   return jnp.where(x > radius, 0., out)
 
 
@@ -63,10 +62,10 @@ def compute_weight_mat(input_size: int, output_size: int, scale,
   weights = kernel(x)
 
   total_weight_sum = jnp.sum(weights, axis=0, keepdims=True)
-  with np.errstate(invalid='ignore', divide='ignore'):
-    weights = jnp.where(
-        jnp.abs(total_weight_sum) > 1000. * np.finfo(np.float32).eps,
-        weights / total_weight_sum, 0)
+  weights = jnp.where(
+      jnp.abs(total_weight_sum) > 1000. * np.finfo(np.float32).eps,
+      jnp.divide(weights, jnp.where(total_weight_sum != 0,  total_weight_sum, 1)),
+      0)
   # Zero out weights where the sample location is completely outside the input
   # range.
   # Note sample_f has already had the 0.5 removed, hence the weird range below.
diff --git a/tests/image_test.py b/tests/image_test.py
index 430c49955..acf98ced6 100644
--- a/tests/image_test.py
+++ b/tests/image_test.py
@@ -324,6 +324,41 @@ class ImageTest(jtu.JaxTestCase):
     output = jax.jit(jit_fn)(x, scale_a, translation_a)
     self.assertAllClose(output, expected, atol=2e-03)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "antialias={}".format(antialias),
+       "antialias": antialias}
+      for antialias in [True, False]))
+  def testScaleAndTranslateGradFinite(self, antialias):
+    image_shape = [1, 6, 7, 1]
+    target_shape = [1, 3, 3, 1]
+
+    data = [
+        51, 38, 32, 89, 41, 21, 97, 51, 33, 87, 89, 34, 21, 97, 43, 25, 25, 92,
+        41, 11, 84, 11, 55, 111, 23, 99, 50, 83, 13, 92, 52, 43, 90, 43, 14, 89,
+        71, 32, 23, 23, 35, 93
+    ]
+
+    x = jnp.array(data, dtype=jnp.float32).reshape(image_shape)
+    scale_a = jnp.array([1.0, 0.35, 0.4, 1.0], dtype=jnp.float32)
+    translation_a = jnp.array([0.0, 0.2, 0.1, 0.0], dtype=jnp.float32)
+
+    def scale_fn(s):
+      return jnp.sum(jax.image.scale_and_translate(
+        x, target_shape, (0, 1, 2, 3), s, translation_a, "linear", antialias,
+        precision=jax.lax.Precision.HIGHEST))
+
+    scale_out = jax.grad(scale_fn)(scale_a)
+    self.assertTrue(jnp.all(jnp.isfinite(scale_out)))
+
+    def translate_fn(t):
+      return jnp.sum(jax.image.scale_and_translate(
+        x, target_shape, (0, 1, 2, 3), scale_a, t, "linear", antialias,
+        precision=jax.lax.Precision.HIGHEST))
+
+    translate_out = jax.grad(translate_fn)(translation_a)
+    self.assertTrue(jnp.all(jnp.isfinite(translate_out)))
+
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())