Fixed error in grad_chooser (for e.g., max) when dtype is not numpy float64

autograd/numpy/numpy_grads.py

@@ -249,8 +249,12 @@ def grad_chooser(g, ans, vs, gvs, x, axis=None, keepdims=None):
     """Builds gradient of functions that choose a single item, such as min or max."""
     g_repeated, _ = repeat_to_match_shape(g, vs, axis, keepdims)
     argmax_locations = x == repeat_to_match_shape(ans, vs, axis, keepdims)[0]
+    if onp.isscalar(x.value):
+        dt = onp.array(x.value).dtype
+    else:
+        dt = x.dtype
     return g_repeated * argmax_locations \
-        / onp.sum(argmax_locations, axis=axis, keepdims=True)
+                      / onp.sum(argmax_locations, axis=axis, keepdims=True).astype(dt)
 
 anp.max.defvjp(grad_chooser)
 anp.min.defvjp(grad_chooser)

autograd/util.py

@@ -22,6 +22,8 @@ def unary_nd(f, x, eps=EPS):
     vs = vspace(x)
     nd_grad = np.zeros(vs.size)
     x_flat = vs.flatten(x)
+    if x_flat.dtype != np.float64:
+        nd_grad = nd_grad.astype(x_flat.dtype)
     for d in range(vs.size):
         dx = np.zeros(vs.size)
         dx[d] = eps/2

tests/test_numpy.py

@@ -81,6 +81,16 @@ def fun(x): return to_scalar(np.max(x))
     check_grads(fun, mat)
     check_grads(d_fun, mat)
 
+def test_max_dtype():
+    """Tests that a dtype other than float64 does not throw an error
+    with the gradient of max.
+    """
+    def fun(x): return to_scalar(np.max(x, 1))
+    d_fun = lambda x : to_scalar(grad(fun)(x))
+    mat = npr.randn(10, 11).astype(np.float32)
+    check_grads(fun, mat)
+    check_grads(d_fun, mat)
+
 def test_max_axis():
     def fun(x): return to_scalar(np.max(x, axis=1))
     d_fun = lambda x : to_scalar(grad(fun)(x))