Is there a function (or possible) to get the inverse of a JAX function?

[nasyxx]

We have the auto grad of a function; can we have the auto inverse of a function?

In this way, when we want to calculate the inverse function of y = f(x) = sin (x ** 2 - 1 ), we don't need to manually write x = f^(-1)(x) = sqrt(arcsin(y) + 1), just an inverse(f)(y).

[jakevdp]

JAX did have an experimental inverse transform at one point, but it was never fully implemented and has been removed.

If you want to see how it might look, the custom jaxpr interpreter tutorial uses an invert transformation as its example application.

[mattjj]

I think there's an oryx.core.inverse function in Oryx; see Automatic function inversion in this colab. (Developing that jaxpr interpreter tutorial is part of what led to Oryx!)

[nasyxx]

If you want to see how it might look, the custom jaxpr interpreter tutorial uses an invert transformation as its example application.

Thanks! Looks cool!

[nasyxx]

I think there's an oryx.core.inverse function in Oryx;

Nice! I never expected such a function in such a probability library!

[krzysztofrusek]

Distrax also has an inv transform https://github.com/deepmind/distrax/blob/b4d78b17ae4aaa3d25c8ca32430dfac8f826d299/distrax/_src/utils/transformations.py#L112