Gradient of pytorch function in jax

[jmsull]

Hi, sorry if this is a naive question, but I can't seem to find much about it anywhere (though there is some related discussion on blackjax).

I have a pytorch function that I wrapped to convert its output to jax arrays, and have defined a custom vjp for it that just wraps the pytorch autodiff gradient.

Calling this function or its gradient is fine, but when trying to sample it in blackjax I find:

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=1/2)>
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

The reason I am asking here is that I have more of a JAX question that might help me understand what is wrong with this- what is the meaning of the "level" in Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=1/2)>?

[jakevdp]

The level is not particularly relevant here – the issue is that your function is being called within JIT-compiled code, and so you're attempting to convert a traced variable to a numpy array, and this is not possible. You can get the same error by doing something like this:

import jax
import numpy as np
import jax.numpy as jnp

@jax.jit
def call_numpy_sin(x):
  x_np = np.asarray(x)
  out_np = np.sin(x_np)
  return jnp.asarray(out_np)

call_numpy_sin(jnp.arange(10))

To understand what's going on here, you might read through How To Think In JAX, in particular the JIT Mechanics section.

The short version of the issue is, you can't convert a traced variable to a numpy array.

I don't know of any good example of interoperability between JAX and pytorch within JIT-compiled code, but you may be able to hack it together using the jax.pure_callback mechanism (it's not very well documented yet, but you can find a short example in #12180 (comment)). Note however that if your'e using this on GPU, it would mean that every call to the pytorch function would require a round trip through the host, so it probably wouldn't be very performant.

Hope that helps!

[jmsull]

Ok thanks for the explanation - I'll try the jax.pure_callback analogous to that comment, for now I am not especially worried about the performance. Thanks!

[jmsull]

This runs through the blackjax NUTS routine now, but takes over an hour for 4k samples (takes ~10 secs for metropolis hastings for the same number of samples), so I suppose the performance is an issue. I will re-write my torch code in JAX, but would be interested to see an example for external logpdfs in the future with blackjax NUTS as discussed in blackjax-devs/blackjax#136.