[Python] How is the execution of Jax and non-Jax parts interleaved in a Python program and...

I have the following code:

def non_jitted_setup():
print("This code runs once at the beginning of the program.")
return jnp.array([1.0, 2.0, 3.0])

class A:

@partial(jax.jit, static_argnums=0)
def my_jitted_function(self, x):
print("This code runs once during the first trace.")
y = x * 2
self.temp = y
return y

# Program execution
data = non_jitted_setup()
A = A()
result1 = A.my_jitted_function(data) # Tracing happens here

np.array(result1)
np.array(A.temp)


If I understand correctly, Jax runs the program line by line and traces the jitted function and runs the Python code inside it once whenever it needs to be compiled and uses the cached version otherwise.

Once y is is returned into result1 above, result1 becomes concrete and can be converted to a numpy.array. However, A.temp still seems to an abstract array despite it being assigned y which is what was returned and concretised to result1 in the previous line, because I get the following error when trying to convert it:

jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[3]


When will the value in A.temp be made concrete? Can we make the value in A.temp be concrete somehow as we know it is used outside the jitted function after it is called?

Continue reading...