Using a HashableArrayWrapper inside JIT'ted function fails

[srijaniiserprinceton]

I am still in the process of learning how to think in JAX. I have been using a HashableArrayWrapper to pass arrays as static arguments to functions along the lines of this answer. This was working for me quite well. But for some reason that I don't quite understand, it has been failing in this script. This collab script should be accessible by anyone.

I decorate the function calc_rad_from_coor with gnool_jit (that lets me pass hashable arrays as static arguments) which is called from within the function model.

@partial(gnool_jit, static_array_argnums=(0,1))
def calc_rad_from_coor(x_m, y_m):
    return np.sqrt(x_m**2 + y_m**2)

def model(r_model):
    npts = len(theta_d)
    
    x_model = r_model * np.cos(theta_d)
    y_model = r_model * np.sin(theta_d)
    
    # adding gjit function here                                                                
    r_model_gjit = calc_rad_from_coor(x_model, y_model)
    
    mu = r_model - r_d
    return np.sum(-0.5*(mu/sigma_d)**2)

Note that model is called from within another jax.jitted function rw_metropolis_kernel.

I have, in the past used gnool_jitted functions within jax.jitted functions but it had worked. Not sure what is happening here. :(

[jakevdp]

I don't entirely understand the motivation for your gnool_jit, but within a jit-compiled function, the result of something like x_model = r_model * np.cos(theta_d) is not an array, it is a tracer. When you pass this tracer to a function that expects it to be a static array, the result will probably not work as expected.

[srijaniiserprinceton]

@jakevdp yep, that's precisely the culprit I think. As I mentioned above, we removed all the partial decorators for the functions where we pass the static arguments but make sure that we only do usual python and numpy ops on the static arguments so that they are never converted to dynamic traced values. This causes us to retain a lot of np operation and usual pythonic for loops.

The reason we were forced to do this is because somehow when calling a jit-decorated function within model (which is the outermost jitted function), the returned values are traced (which we don't expect since they depend on static arguments).

It would be great you could this Collab notebook which contains a simplified example of the overall structure of the code that we want to work and let us know how to make the very last call to model_jnp_(GVARS) functional (which, as you will see fails since the arguments that are getting passed to a subsequent function is getting traced).

[jakevdp]

Thanks. One comment on a read-through of the notebook. In these lines:

        # calling the function with two static arguments
        a_static, b_static = func1_w_staticargs_jnp(GVARS.a, GVARS.b)

        # when jitted, a_static and b_static become traced here
        # even thought we declare the argument to be static and use numpy ops only
        print('a_static, b_static: \n', a_static, b_static)

The function func1_w_staticargs_jnp creates DeviceArrays, and any device array created within a JIT context will be traced. From the comment, it seems like maybe you meant to call func1_w_staticargs_np? If that were the case, I think the outputs would still be static.

[srijaniiserprinceton]

What I did in the collab notebook was to demonstrate why we are using numpy ops in our code which is making the compilation time long. So, it was intentional to call func1_w_staticargs_jnp. We discussed in your previous reply how non-jax operations (loops etc) are rolled out and then square of the number of operations are performed. So, in the notebook I was trying to show you why we are forced to use the numpy ops (in func1_w_staticargs_np) and why using the jax.numpy counterpart (in func1_w_staticargs_jnp) fails. It will be possible to stop the rolling out of loops in our actual code, if a function like func1_w_static_jnp which is decorated with static_argnums returns static objects (in which case we can replace the numpy counterpart with the jax.numpy counterpart.

So, in short:

• We currently use numpy ops on static arguments -> Makes compilation time long. (So what you wrote in your suggestion in the above reply is what we are doing now in the actual code.
• We are forced to do the above since the jax.numpy ops returns us traced objects which cannot be passed into a subsequent function that wants to use these values for creating shaped matrices.

[srijaniiserprinceton]

One of the options that I can think of to bypass this passing of static arguments is described below:

• I fix the shapes of the matrices that I need to make to the largest matrix that I would ever encounter. Let's call this large matrix M_large.
• I pass the nestedtuples (or arrays) which are traced (not static) to slice out and fill in smaller areas in the large matrix M_large. Is this possible by dynamic_slice?

If it is possible to fill out smaller patches in M_large (the size of these small patches being governed by a traced value), then we can perhaps exclude static arguments altogether and make everything numpy instead of jax.numpy.

[jakevdp]

Oh, sorry, I got distracted by the comment, which says that np ops are being used, when it is jnp ops that are being used. Maybe it's just outdated.