Using a HashableArrayWrapper inside JIT'ted function fails

[srijaniiserprinceton]

This is the same as the question I posted in the JAX discussion forum. But since the source code is the same as that found in this script on JAX sampler by @rlouf , I thought it might be more relevant to post it here as well.

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. :(