Possible bug for jax.numpy.abs complex-input grad

Please:

• Check for duplicate issues. (tried my best to find a duplicate on this)
• Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

In [1]: import jax; import jax.numpy as jnp
In [2]: t1 = jnp.array([[1.0, 2.0], [0.0, -1.0]])
In [3]: jax.grad(lambda t: jnp.sum(jnp.abs(t)))(t1)
Out[3]: 
DeviceArray([[ 1.,  1.],
             [ 1., -1.]], dtype=float32)
In [7]: t2 = t1 * 1+0j # this is so we convert to complex64
In [8]: jax.grad(lambda t: jnp.sum(jnp.abs(t)))(t2)
Out[8]: 
DeviceArray([[ 1.+0.j,  1.+0.j],
             [ 0.+0.j, -1.+0.j]], dtype=complex64)

In the code sequence above, we can compare the grad of basically the same input as a float tensor and as a complex tensor. The problem I have (and for which I could not find any references) is that when an entry is 0, the grad changes from 1 (at float input) to 0 (at complex input), even though we’re talking about the same abs function applied to said entry.

From looking at the source code, this seems like a matter of using select to replace the result at the positions which were originally zero with 1, after the calculation. I’d be happy to submit a PR if this is indeed a bug. Otherwise, perhaps a reference on why this is intended would be great!

[mattjj]

Thanks for the question!

As you’ve already discussed, the observation which started this issue comes down to this:

import jax
import jax.numpy as jnp

print(jax.grad(jnp.abs)(0.))  # 1.0
print(jax.grad(jnp.abs)(0. + 0.j))  # 0j

Also as already discussed, the mathematical function being represented here (complex absolute value and its restriction to real numbers) isn’t differentiable at the point we’re evaluating the gradient. So we could return a nan here to represent the non-differentiability. But instead we usually just provide some implementation-defined answer which seems expedient. Yet our implementation-defined answer differs in the real- and complex-input cases, not for deep reasons but just an artifact of how we happened to write the abs JVP rule.

I think there are four options:

1. make grad(abs)(0.) and grad(abs)(0.+0.j) both be nan
2. make grad(abs)(0.) == grad(abs)(0.+0.j) == 1. (currently what we do for real-valued inputs only)
3. make grad(abs)(0.) == grad(abs)(0.+0.j) == 0. (currently what we do for complex-valued inputs only)
4. do nothing, keeping grad(abs)(0.) == 1. and grad(abs)(0.+0.j) == 0..

I bet if we choose option 1 we’ll break existing code (which I’m guessing relies on not producing nans here). If we choose option 3 we’d be changing the behavior for real-valued inputs, which I suspect is much more common.

Both options 2 and 4 seem viable. Let’s try 2 and see if anyone complains!

[polvalente]

@YouJiacheng exactly that. The issue is that the same number is differentiating differently even though only the datatype changed, not the number itself.

I just continue to be unsure if this is a decision or just a side-effect of how it was implemented, which is the main reason on why I opened this issue.

re the edit above: I get that the complex extension opens up more paths for the derivative to be checked. If efficiency is the sole reason why they are different, perhaps a comment on the jvp code might suffice.