Advice on speeding up compilation time

[isaacgerg]

I am converting code I wrote in Ceres to Jax so that i can debug and tune it easier. The code is very straightforward but takes quite a bit of time to compile. I am looking for advice on how to speed this up. Perhaps I can reformulate my code or use a feature I am not aware of and this will save me some time. Any help appreciated.

I am running jax 0.1.72 and jaxlib 0.1.51. I am compiling on for CPU.

def motion_model_wrapper(x0, delays, t_ref, weights, v_1, v_2, p1_attitude, p2_attitude, p_tx, p1_rx, p2_rx, c, altitude, p_dt):
   # i am using the scipy optimize lib so i have to enable jax.config.update('jax_enable_x64', True) before calling .jit()
    # x0 is 4 x 1
    # delays, t_ref, weights,  is 100 x 1
    # v_1, v_2, c, altitude, p_dt are scalars
    # p1_attitude, p2_attitude, p_tx, p1_rx, p2_rx are 3 x 1
    num_measurements = len(delays)  # This is a always 100.  
    #e = jnp.zeros(num_measurements)
    vy_1 = x0[0]; vz_1 = x0[1]; 
    vy_2 = x0[0]; vz_2 = x0[1];   # TODO remove in the future
    
    total_error = 0
    for k in range(num_measurements):
        p1_tx = rotate3d(p_tx, p1_attitude) + jnp.array([0, 0, -altitude])
        p1_rx = rotate3d(p1_rx, p1_attitude) + jnp.array([v_1[0]*t_ref[k], vy_1*t_ref[k], -altitude + vz_1*t_ref[k]])
        p2_tx = rotate3d(p_tx, p2_attitude) + jnp.array([v_1[0]*p_dt, 0, -altitude])
        p2_rx = rotate3d(p2_rx, p2_attitude) + jnp.array([v_1[0]*p_dt, vy_1*p_dt , vz_1*p_dt]) + jnp.array([v_2[0]*t_ref[k], vy_2*t_ref[k], -altitude + vz_2*t_ref[k]])        
        r = motion_model(p1_tx, p2_tx, p1_rx, p2_rx, delays[k], t_ref[k], c, altitude)
        total_error = total_error + weights[k] * jnp.abs(r)
        #e[k] = motion_model(p1_tx, p2_tx, p1_rx, p2_rx, delays[k], t_ref[k], c, altitude)
        #jax.ops.index_update(e, jax.ops.index(k), motion_model(p1_tx, p2_tx, p1_rx, p2_rx, delays[k], t_ref[k], c, altitude))

    #print(total_error)

    return total_error

[hawkinsp]

The key to faster compilation time is to avoid the Python for loop, which unrolls the loop and creates a very large computation.

Two options you might consider:

• if each measurement is independent, try using jax.vmap.
• another alternative is to use a JAX-level control flow construct, such as lax.fori_loop, lax.map or lax.scan.

[isaacgerg]

Aside from the summing of the error, each iteration is independent. But, you could do each iteration interdependently and store the error to a vector which is summed for the final error. It's not clear to me how one would apply vmap here from looking at the examples. Specifically, its not clear to me how I would map k into body_fun() so that vmap could be used.

[hawkinsp]

You don't necessarily need to pass k in, given it seems you are using it to index into various arrays. You should simply use the in_axes specification to vmap to map over those dimensions.

That said, if you really did want to pass k in, simply vmap over jnp.arange(num_measurements).

[isaacgerg]

I think I figured it out. Sort off...

I packaged up the inputs into an dictionary and keep tally of my error.

    val = {'x0':x0, 'delays':delays, 't_ref':t_ref, 'weights':weights, 'v_1':v_1, 'v_2':v_2, 'p1_attitude':p1_attitude, 
           'p2_attitude':p2_attitude, 'p_tx':p_tx, 'p1_rx':p1_rx, 'p2_rx':p2_rx, 'c':c, 'altitude':altitude, 'p_dt':p_dt, 'error':0.0}
    val = jax.lax.fori_loop(0, num_measurements, fun, val)

However, I am calling grad on this function since I am using jax to estimate the jacobian for use in scipy.optimize. Jax is giving me this now.

builtins.ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead.

Am i using for_i correctly?

I will try scan now.

[isaacgerg]

I got it working by using .scan(). I made carry = a dictionary of my input parameters. I stored the error into y. I just used carry as a pass through of the params by having my function return carry, error. I made the input to scan for xs = np.arange(num_measurements).