Add custom_jvp to jax.numpy.ldexp.

[carlosgmartin]

Addresses #11467 (comment).

[jakevdp]

Thanks for looking into this! I'm not sure this is actually the correct gradient behavior for this function. As I mentioned in #11467 (comment), ldexp is a bit of a strange function, because it sort of computes x * 2 ** y, but doesn't really compute that, because it's operating on the details of the bitwise representation of x and y. Because of that, I think autodiff doesn't really apply to this function, because autodiff operates only in the space of real numbers being represented by the floating point implementation, so bitwise manipulations don't really have a gradient.

Maybe the best solution here would be to define custom_jvp that raises an error saying basically "this function isn't differentiable". What do you think?

[carlosgmartin]

@jakevdp Thanks for your feedback.

It seems to me that, if a JAX function extensionally computes a differentiable mathematical function $f$ then, regardless of how it is implemented internally, it ought to have the same gradient.

Indeed, my understanding is that one of the primary use cases of custom_jvp is precisely to recover gradients of functions whose internal implementations are not auto-differentiable.

Since ldexp computes $f(x) = x 2^n$, the least surprising behavior from a user POV would be that $f'(x) = 2^n$. Indeed, a user (or the compiler) might encounter such an expression in a program and substitute it with ldexp as an optimization. It would seem surprising if this substitution were to cause an error on a backward pass.

Thoughts?

[jakevdp]

So my question is, if we're just computing $x * 2 ^ y$, why don't we fix this by changing the implementation to

def ldexp(x, y):
  return x * 2 ** y

Then no custom JVP is necessary at all.

[carlosgmartin]

I think (the bit-twiddling implementation of) ldexp is supposed to be a faster way to do that, at least on some platforms.

The CUDA Math API has dedicated ldexp functions for single and double precision.

The C standard library also has dedicated ldexp functions.

Not sure what the performance advantage is for different platforms.

[jakevdp]

Sure, but JAX does not dispatch to any of those fast kernels, and I imagine the current bit-twiddling implementation is far slower than just writing x1 * 2 ** x2. Is there any good reason to keep the current implementation?

[jakevdp]

I guess stepping back, here are the options:

1. ldexp is fundamentally a bit-twiddling operation. In this case, its autodiff behavior is poorly defined, and if we add a custom_vjp, it should probably just raise an error.
2. ldexp represents the mathematical operation $x * 2^y$ with platform-specific implementation details when custom kernels are available. In this case, the optimal implementation in JAX would be to write x * 2 ** y, which would have the side effect of making custom_jvp unnecessary.

Until now, we've approached this as (1). Which do you think is the right approach?

[carlosgmartin]

You raise an interesting point. If there's indeed no performance advantage to the bit-twiddling implementation of ldexp(x, n) over x * 2 ** n, then perhaps the former isn't necessary and can be replaced with the latter.

At least for the time being, perhaps it's worth adding a note to the documentation for ldexp stating that there's no performance advantage to its current bit-twiddling implementation over x * 2 ** n.

I'd also welcome any additional opinions from people who are more familiar with the hardware side of things.

[jakevdp]

Looks good!

One other thing: we should update the function docs with info about the implementation (this would involve removing the @implements decorator and writing a full docstring).

If you'd like to do this as part of the PR then go ahead, but I'm happy to update docs in a followup.

[carlosgmartin]

I'll let you handle that so you can choose the best wording.

[jakevdp]

Thanks for putting this together!