Code runs half as fast after restarting jupyter notebook the next day

[alexandermm]

Hi, I have some code using a custom autodiff function used by another function pSpatialComp. Just using CPU. Then I find the Jacobian and Hessian of the function for different parameters in a loop. For some reason yesterday this was taking 40 seconds and now it takes >80 seconds.

What could the issue be? Is it a memory management issue? How could I make the code faster? The code is shown below:

import jax
import jax.numpy as jnp
import scipy.special

import time

jax.config.update("jax_enable_x64", True)

#Types being used
realType = jnp.float64
compType = jnp.complex64


@jax.custom_jvp
def modJv(v,a,z):
    v, a, z = jnp.asarray(v), jnp.asarray(a), jnp.asarray(z)

    # Require the order v to be integer type: this simplifies the JVP rule below
    assert jnp.issubdtype(v.dtype, jnp.integer)
    # Promote the input to inexact (float/complex), jnp.result_type() accounts for enable_x64 flag
    z = z.astype(jnp.result_type(float, z.dtype))
    a = a.astype(jnp.result_type(float, a.dtype))

    # Wrap scipy function to return the expected dtype
    a_sqrtz = a * jnp.sqrt(z)
    def _scipy_jv(v, a_sqrtz):
        return scipy.special.jv(v, a_sqrtz).astype(z.dtype)

    # Define the expected shape & dtype of output
    result_shape_dtype = jax.ShapeDtypeStruct(
        shape=jnp.broadcast_shapes(v.shape, a.shape, z.shape), dtype=z.dtype)

    # We use vectorize=True because scipy.special.jv handles broadcasted inputs
    return jax.pure_callback(_scipy_jv, result_shape_dtype, v, a_sqrtz, vectorized=True)


@modJv.defjvp
def _modJv_jvp(primals, tangents):
    v, a, z     = primals
    _, _, z_dot = tangents
    
    Jv, Jvm2, Jvp2 = modJv(v,a, z), modJv(v-2,a, z), modJv(v+2,a, z)

    dmodJv_dz = a*a/8.0 * ((Jv+Jvm2)/(v-1) - (Jv+Jvp2)/(v+1))
    return modJv(v,a, z), z_dot * dmodJv_dz


@jax.jit
def pSpatialComp(point, transLocation,transNormal, kr_trans):
    dirV = 1 #Nu parameter for Bessel function in transducer field directivity function
    
    distVec   = (point-transLocation).astype(compType)
    dist2     = jnp.sum(jnp.square(distVec))
    sinTheta2 = jnp.sum(jnp.square(jnp.cross(distVec,transNormal).astype(compType)))/dist2
    
    directivity = 0.25*(modJv(dirV+1,kr_trans, sinTheta2) + modJv(dirV-1,kr_trans, sinTheta2))
    
    dist = jnp.sqrt(dist2)
    
    return directivity/dist*jnp.exp(-1j*k*dist)


c_m     = 332.0
f_trans = 40.0e3

k        = 2*jnp.pi*f_trans/c_m
r_trans  = 5.0e-3
kr_trans = k*r_trans


num_trans_x = 10 #Should be even
num_trans_y = 10 #Should be even
num_trans   = num_trans_x * num_trans_y
trans_dx    = 10.0e-3
trans_z_pos = 0.0e-3

plane_z_pos      = 60.0e-3
plane_opt_radius = 10.0e-3
plane_num_ps     = 60


indices = jnp.reshape(jnp.arange(0, plane_num_ps, dtype=realType) + 0.5, (plane_num_ps,1))
r_ps    = plane_opt_radius * jnp.sqrt(indices/plane_num_ps)
angle   = jnp.pi * (1.0 + jnp.sqrt(5.0)) * indices

plane_ps_x = r_ps * jnp.cos(angle)
plane_ps_y = r_ps * jnp.sin(angle)
array_z    = jnp.full((plane_num_ps,1), plane_z_pos)

plane_ps_imag = jnp.concatenate((plane_ps_x, plane_ps_y, array_z), axis=1, dtype=compType)


x_vals = jnp.linspace(-(num_trans_x/2-0.5)*trans_dx, (num_trans_x/2-0.5)*trans_dx, num_trans_x)
y_vals = jnp.linspace(-(num_trans_y/2-0.5)*trans_dx, (num_trans_y/2-0.5)*trans_dx, num_trans_y)
trans_x,trans_y = jnp.meshgrid(x_vals, y_vals)
trans_x         = jnp.reshape(trans_x, (-1,1))
trans_y         = jnp.reshape(trans_y, (-1,1))
trans_z         = jnp.full((num_trans,1), trans_z_pos)
transLocations  = jnp.concatenate((trans_x, trans_y, trans_z), axis=1, dtype=realType)

transNormals    = jnp.tile(jnp.array([0.0, 0.0, 1.0],dtype=realType), (num_trans,1))


g = lambda point: pSpatialComp(point, transLocation, transNormal, kr_trans)

pxVals = jnp.empty(shape=(plane_num_ps, num_trans),      dtype=compType)
pxJacs = jnp.empty(shape=(plane_num_ps, num_trans, 3),   dtype=compType)
pxHess = jnp.empty(shape=(plane_num_ps, num_trans, 3,3), dtype=compType)

start_time = time.time()
for p in range(plane_num_ps):
    for t in range(num_trans):
        point         = plane_ps_imag[p,:]
        transLocation = transLocations[t,:]
        transNormal   = transNormals[t,:]
        
        pxVals = pxVals.at[p,t].set(g(point))
        
        J = jax.jacfwd(g, holomorphic=True)
        pxJacs = pxJacs.at[p,t,:].set(J(point))
        
        pxHess = pxHess.at[p,t,:,:].set(jax.jacfwd(J, holomorphic=True)(point))
hess_time  = time.time() - start_time
print("Hessian comp FF (complex, assuming holomorphic): %.6f secs" % hess_time)
print("\n")

print(pxVals[1:2,1:2])
print("\n")
print(pxJacs[1:2,1:2,:])
print("\n")
print(pxHess[1:2,1:2,:,:])
print("\n")

[jakevdp]

Hard to say, but one guess: if you set jax.config.update("jax_enable_x64", True) after executing some code, there could be cached 32-bit operations that will execute on subsequent runs. So perhaps yesterday you ran some code before setting the flag, and you were executing faster 32-bit operations.

[jakevdp]

Any time you're writing for-loops over array axes in numpy or JAX it will generally lead to poorer performance. Have you tried vmap?

[alexandermm]

Hi! Yes I've tried it but haven't been able to get it working. I tried doing it for one of the loops, however the shape of the output array don't seem to be correct. What is the issue here?

Also, is there a way of getting the Jacobian information from the Hessian? Those values should be in there when I use jit for jax.jacfwd(jax.jacfwr(g))

The code is shown below:

g = jax.jit(pSpatialComp)
J = jax.jit(jax.jacfwd(g, argnums=0, holomorphic=True))
H = jax.jit(jax.jacfwd(J, argnums=0, holomorphic=True))


pxVals = jnp.empty(shape=(plane_num_ps, num_trans),      dtype=compType)
pxVals2 = jnp.empty(shape=(plane_num_ps, num_trans),      dtype=compType)

pxJacs = jnp.empty(shape=(plane_num_ps, num_trans, 3),   dtype=compType)
pxJacs2 = jnp.empty(shape=(plane_num_ps, num_trans, 3),   dtype=compType)

pxHess = jnp.empty(shape=(plane_num_ps, num_trans, 3,3), dtype=compType)
pxHess2 = jnp.empty(shape=(plane_num_ps, num_trans, 3,3), dtype=compType)


start_time = time.time()
for p in range(plane_num_ps):
    for t in range(num_trans):
        point         = plane_ps_imag[p,:]
        transLocation = transLocations[t,:]
        transNormal   = transNormals[t,:]
        
        pxVals = pxVals.at[p,t].set(g(point, transLocation, transNormal, kr_trans))       
        pxJacs = pxJacs.at[p,t,:].set(J(point, transLocation, transNormal, kr_trans))
        pxHess = pxHess.at[p,t,:,:].set(H(point, transLocation, transNormal, kr_trans))
hess_time  = time.time() - start_time
print("Hessian fwd(fwd) diff types (complex, assuming holomorphic): %.6f secs" % hess_time)
print("\n")


g_v = jax.vmap(g, in_axes=(None, 0, 0, None), out_axes=0)
J_v = jax.vmap(J, in_axes=(None, 0, 0, None), out_axes=0)
H_v = jax.vmap(H, in_axes=(None, 0, 0, None), out_axes=0)

start_time = time.time()
for p in range(plane_num_ps):
    point  = plane_ps_imag[p,:]
    #point  = jnp.expand_dims(point, 0)
        
    output = g_v(point, transLocations, transNormals, kr_trans)
    pxVals2 = pxVals2.at[p,:].set(output)
    
    output  = J_v(point, transLocations, transNormals, kr_trans)
    pxJacs2 = pxJacs.at[p,:,:].set(output)
    
    output  = H_v(point, transLocations, transNormals, kr_trans)
    pxHess2 = pxHess2.at[p,:,:,:].set(output)
hess_time  = time.time() - start_time
print("Hessian fwd(fwd) diff types (complex, assuming holomorphic): %.6f secs" % hess_time)
print("\n")

[jakevdp]

If you're trying to replace a nested for loop with vmap, you'll need a nested vmap.

[alexandermm]

Hi Jake, thank you I'm trying to get vamp working in one dimension before trying to use nested vamps; when I use the line g_v = jax.vmap(g, in_axes=(None, 0, 0, None), out_axes=0) I get the right output but the array does not have the shape N by 1, where N is the size of the axes where I'm using zero in in_axes. How is the output shape determined in vmap?

[jakevdp]

The output shape of y = jax.vmap(f, in_axes=0, out_axes=0)(x) will be (x.shape[0], *f(x[0]).shape). That is, x.shape[in_axes] will be the same as y.shape[out_axes], with the rest of y.shape determined by the shape of the output of f for the unbatched inputs. Does that make sense?