Repository: google/jax

Issue:
jax.numpy.tanh(2) works but jax.grad(tanh)(2) fails- int32 not liked by grad
Hi 
this may be a little churlish to bring up but ints are acceptable for some calls in jax but not for others. I assume this is not by design. 
example
`jax.version.__version__ == '0.1.62'`
```
import jax.numpy as np
from jax import grad
aval =2
print("aval=" +str(aval))
print("np.tanh(aval)=" + str(np.tanh(aval)))

try:
    grad(np.tanh)(aval)
except TypeError as e:
    print(e)
    print ("sadness - grad(np.tanh)(aval) failed ")
```
Output
aval=2
np.tanh(aval)=0.9640276
Primal inputs to reverse-mode differentiation must be of float or complex type, got type int32
sadness - grad(np.tanh)(aval) failed 







Discussion:
Well, I think this particular example is working as intended. You can think of your original `tanh` function as equivalent to the following:

```
import jax.numpy as jnp

def f(x):
  x = jnp.array(x, jnp.float32)
  return jnp.tanh(x)
```

In other words, the int->float cast happens *inside* the function you are differentiating. That is, you are differentiating a function that takes as input an integer and returns a float.

In general the cotangent types returned by `grad` match the input types, so we'd have to return an integer value to you, which probably isn't what you wanted. An error seems preferable.

What do you think?
Hi Peter,

Thank you for the quick response.

I think your comment is quite reasonable and I get that  the cotangent code is driven by the function and its inputs so the results  are understandable.

I can think of at least two ways forward.

1 Take the expedient approach and just make it clear (may be more clear if it is already in the docs; I apologize if this issue is discussed in the docs and I missed it) that the input value controls the output value from the cotangent and that since ints for gradients are not sensible that the inputs should be floats otherwise the gradient code will throw an error.

2 Take a potentially more hazardous approach and handle ints in a more graceful way by casting them to a default float type before really starting on making the cotangent

Given that there are probably better things to spend one's time on and approach 2 may give rise to some nasty subtle issues,  I would think that option 1 would be the more sensible way of moving forward.

Dominic


________________________________
From: Peter Hawkins <notifications@github.com>
Sent: Thursday, April 9, 2020 3:06 PM
To: google/jax <jax@noreply.github.com>
Cc: Barraclough, Dominic (ext. 414) <djb@qvii.com>; Author <author@noreply.github.com>
Subject: [EXTERNAL] Re: [google/jax] jax.numpy.tanh(2) works but jax.grad(tanh)(2) fails- int32 not liked by grad (#2663)

CAUTION:
This email originated from outside of QVI. Do not click links or open attachments unless you recognize the sender and are expecting to receive this content from them. If you suspect the email is not legitimate, please forward it as an attachment to SPAM@qvii.com and delete it from your Inbox.


Well, I think this particular example is working as intended. You can think of your original tanh function as equivalent to the following:

import jax.numpy as jnp

def f(x):
  x = jnp.array(x, jnp.float32)
  return jnp.tanh(x)


In other words, the int->float cast happens inside the function you are differentiating. That is, you are differentiating a function that takes as input an integer and returns a float.

In general the cotangent types returned by grad match the input types, so we'd have to return an integer value to you, which probably isn't what you wanted. An error seems preferable.

What do you think?

—
You are receiving this because you authored the thread.
Reply to this email directly, view it on GitHub<https://github.com/google/jax/issues/2663#issuecomment-611701532>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AAIRK64NOX52JSGRAM4MD53RLYMC7ANCNFSM4ME5ESRQ>.


Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -334,7 +334,9 @@ def grad(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
 
   Args:
     fun: Function to be differentiated. Its arguments at positions specified by
-      ``argnums`` should be arrays, scalars, or standard Python containers. It
+      ``argnums`` should be arrays, scalars, or standard Python containers.
+      Argument arrays in the positions specified by ``argnums`` must be of
+      inexact (i.e., floating-point or complex) type. It
       should return a scalar (which includes arrays with shape ``()`` but not
       arrays with shape ``(1,)`` etc.)
     argnums: Optional, integer or sequence of integers. Specifies which


Test patch: