jit compile too slow in my code

[XDongiang]

I read the gitbook, I think my problem my be it :
Sometimes it isn’t obvious how to rewrite your code to avoid Python loops because your code makes use of many arrays with different shapes. The recommended solution in this case is to make use of functions like jax.numpy.where() to do your computation on padded arrays with fixed shape. The JAX team is exploring a “masking” transformation to make such code easier to write.
my code :

def likelihood(self, args):
        
        phi_mass = args[0:1]
        
        phi_width = args[1:2]
        
        kk_f980_mass = args[2:3]
        
        kk_f980_g_kk = args[3:4]
        
        kk_f980_rg = args[4:5]
        
        kk_f980_const = args[5:7].reshape(-1,2)
        
        kk_f980_theta = args[7:9].reshape(-1,2)
        
        kk_f0_mass = args[9:13]
        
        kk_f0_width = args[13:17]
        
        kk_f0_const = args[17:25].reshape(-1,2)
        
        kk_f0_theta = args[25:33].reshape(-1,2)
        
        kk_f1270_mass = args[33:34]
        
        kk_f1270_width = args[34:35]
        
        kk_f1270_const = args[35:40].reshape(-1,5)
        
        kk_f1270_theta = args[40:45].reshape(-1,5)
        
        kk_f2_mass = args[45:56]
        
        kk_f2_width = args[56:67]
        
        kk_f2_const = args[67:122].reshape(-1,5)
        
        kk_f2_theta = args[122:177].reshape(-1,5)
        
        kk_u_rho_mass = args[177:181]
        
        kk_u_rho_width = args[181:185]
        
        kk_u_rho_const = args[185:189].reshape(-1,1)
        
        kk_u_rho_theta = args[189:193].reshape(-1,1)
        
        kk_u_b_mass = args[193:196]
        
        kk_u_b_width = args[196:199]
        
        kk_u_b_const = args[199:211].reshape(-1,4)
        
        kk_u_b_theta = args[211:223].reshape(-1,4)
        
        pipi_f980_mass = args[223:224]
        
        pipi_f980_g_pipi = args[224:225]
        
        pipi_f980_rg = args[225:226]
        
        pipi_f980_const = args[226:228].reshape(-1,2)
        
        pipi_f980_theta = args[228:230].reshape(-1,2)
        
        pipi_f0_mass = args[230:234]
        
        pipi_f0_width = args[234:238]
        
        pipi_f0_const = args[238:246].reshape(-1,2)
        
        pipi_f0_theta = args[246:254].reshape(-1,2)
        
        pipi_f1270_mass = args[254:255]
        
        pipi_f1270_width = args[255:256]
        
        pipi_f1270_const = args[256:261].reshape(-1,5)
        
        pipi_f1270_theta = args[261:266].reshape(-1,5)
        
        pipi_f2_mass = args[266:277]
        
        pipi_f2_width = args[277:288]
        
        pipi_f2_const = args[288:343].reshape(-1,5)
        
        pipi_f2_theta = args[343:398].reshape(-1,5)
        
        pipi_u_rho_mass = args[398:402]
        
        pipi_u_rho_width = args[402:406]
        
        pipi_u_rho_const = args[406:410].reshape(-1,1)
        
        pipi_u_rho_theta = args[410:414].reshape(-1,1)
        
        pipi_u_b_mass = args[414:416]
        
        pipi_u_b_width = args[416:418]
        
        pipi_u_b_const = args[418:426].reshape(-1,4)
        
        pipi_u_b_theta = args[426:434].reshape(-1,4)

How to replace the "reshape" to a good compile speed

[jakevdp]

I'm not sure whether this will improve compile speed, but your approach of manually flattening and unflattening the parameters seems a bit brittle: you might try using jax.flatten_util.ravel_pytree to automate this. For example:

import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
a = np.arange(3.0)
b = jnp.ones((2, 1))
c = jnp.eye(3)
flattened, unravel = ravel_pytree((a, b, c))

print(flattened)
# [0. 1. 2. 1. 1. 1. 0. 0. 0. 1. 0. 0. 0. 1.]

a, b, c = unravel(flattened))
print(a)
# DeviceArray([0., 1., 2.], dtype=float32)
print(b)
# DeviceArray([[1.],
#              [1.]], dtype=float32),
print(c)
# DeviceArray([[1., 0., 0.],
#              [0., 1., 0.],
#              [0., 0., 1.]], dtype=float32)]

[lucasliunju]

Hi jack, I would like to ask you an question. I'm trying to use jax on tpu. In my code, I try to define the tpu:

from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://[My TPU ip address]:8470"`

I'm not sure whether my definition is right. In addition, the tpu version is tpu_driver_nightly.

Thank you very much!

[jakevdp]

Are you using Colab? If so you can run jax.tools.colab_tpu.setup_tpu() and it should wet everything up for you.

[lucasliunju]

I am using my own tpu. I find I can run the code on cpu. However, when I change to tpu, it occurs segmentation error. I guess my definition is wrong.

Best,
Yong

[lucasliunju]

I try to increase the memory of tpu, it still occurs segmentation error.

[XDongiang]

I'm not sure whether this will improve compile speed, but your approach of manually flattening and unflattening the parameters seems a bit brittle: you might try using jax.flatten_util.ravel_pytree to automate this. For example:

import jax.numpy as jnp
from jax.flatten_util import ravel_pytree
a = np.arange(3.0)
b = jnp.ones((2, 1))
c = jnp.eye(3)
flattened, unravel = ravel_pytree((a, b, c))

print(flattened)
# [0. 1. 2. 1. 1. 1. 0. 0. 0. 1. 0. 0. 0. 1.]

a, b, c = unravel(flattened))
print(a)
# DeviceArray([0., 1., 2.], dtype=float32)
print(b)
# DeviceArray([[1.],
#              [1.]], dtype=float32),
print(c)
# DeviceArray([[1., 0., 0.],
#              [0., 1., 0.],
#              [0., 0., 1.]], dtype=float32)]

I didn't do it manually, i use jinja2 to creat the script.

[XDongiang]

I found that the main source of compile time lies in the complexity of my model, When I decrease and increase the number of models, the compilation time changes linearly.

def likelihood(self, args):
        
        phi_mass = args[0:1]
        
        phi_width = args[1:2]
        
        kk_f980_mass = args[2:3]
        
        kk_f980_g_kk = args[3:4]
        
        kk_f980_rg = args[4:5]
        
        kk_f980_const = args[5:7].reshape(-1,2)
        
        kk_f980_theta = args[7:9].reshape(-1,2)
        
        kk_f0_mass = args[9:13]
        
        kk_f0_width = args[13:17]
        
        kk_f0_const = args[17:25].reshape(-1,2)
        
        kk_f0_theta = args[25:33].reshape(-1,2)
        
        kk_f1270_mass = args[33:34]
        
        kk_f1270_width = args[34:35]
        
        kk_f1270_const = args[35:40].reshape(-1,5)
        
        kk_f1270_theta = args[40:45].reshape(-1,5)
        
        kk_f2_mass = args[45:56]
        
        kk_f2_width = args[56:67]
        
        kk_f2_const = args[67:122].reshape(-1,5)
        
        kk_f2_theta = args[122:177].reshape(-1,5)
        
        kk_u_rho_mass = args[177:181]
        
        kk_u_rho_width = args[181:185]
        
        kk_u_rho_const = args[185:189].reshape(-1,1)
        
        kk_u_rho_theta = args[189:193].reshape(-1,1)
        
        kk_u_b_mass = args[193:196]
        
        kk_u_b_width = args[196:199]
        
        kk_u_b_const = args[199:211].reshape(-1,4)
        
        kk_u_b_theta = args[211:223].reshape(-1,4)
        
        pipi_f980_mass = args[223:224]
        
        pipi_f980_g_pipi = args[224:225]
        
        pipi_f980_rg = args[225:226]
        
        pipi_f980_const = args[226:228].reshape(-1,2)
        
        pipi_f980_theta = args[228:230].reshape(-1,2)
        
        pipi_f0_mass = args[230:234]
        
        pipi_f0_width = args[234:238]
        
        pipi_f0_const = args[238:246].reshape(-1,2)
        
        pipi_f0_theta = args[246:254].reshape(-1,2)
        
        pipi_f1270_mass = args[254:255]
        
        pipi_f1270_width = args[255:256]
        
        pipi_f1270_const = args[256:261].reshape(-1,5)
        
        pipi_f1270_theta = args[261:266].reshape(-1,5)
        
        pipi_f2_mass = args[266:277]
        
        pipi_f2_width = args[277:288]
        
        pipi_f2_const = args[288:343].reshape(-1,5)
        
        pipi_f2_theta = args[343:398].reshape(-1,5)
        
        pipi_u_rho_mass = args[398:402]
        
        pipi_u_rho_width = args[402:406]
        
        pipi_u_rho_const = args[406:410].reshape(-1,1)
        
        pipi_u_rho_theta = args[410:414].reshape(-1,1)
        
        pipi_u_b_mass = args[414:416]
        
        pipi_u_b_width = args[416:418]
        
        pipi_u_b_const = args[418:426].reshape(-1,4)
        
        pipi_u_b_theta = args[426:434].reshape(-1,4)
        
        
        d_phif0_kk_BW_flatte980 = self.calculate_BW_flatte980(phi_mass,phi_width,kk_f980_mass,kk_f980_g_kk,kk_f980_rg,kk_f980_const,kk_f980_theta,self.data_phi_kk,self.data_f_kk,self.data_phif0_kk)

        m_phif0_kk_BW_flatte980 = self.calculate_BW_flatte980(phi_mass,phi_width,kk_f980_mass,kk_f980_g_kk,kk_f980_rg,kk_f980_const,kk_f980_theta,self.mc_phi_kk,self.mc_f_kk,self.mc_phif0_kk)
        
        d_phif0_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_f0_mass,kk_f0_width,kk_f0_const,kk_f0_theta,self.data_phi_kk,self.data_f_kk,self.data_phif0_kk)

        m_phif0_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_f0_mass,kk_f0_width,kk_f0_const,kk_f0_theta,self.mc_phi_kk,self.mc_f_kk,self.mc_phif0_kk)
        
        d_phif2_kk_BW_flatte1270 = self.calculate_BW_flatte1270(phi_mass,phi_width,kk_f1270_mass,kk_f1270_width,kk_f1270_const,kk_f1270_theta,self.data_phi_kk,self.data_f_kk,self.data_phif2_kk)

        m_phif2_kk_BW_flatte1270 = self.calculate_BW_flatte1270(phi_mass,phi_width,kk_f1270_mass,kk_f1270_width,kk_f1270_const,kk_f1270_theta,self.mc_phi_kk,self.mc_f_kk,self.mc_phif2_kk)
        
        d_phif2_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_f2_mass,kk_f2_width,kk_f2_const,kk_f2_theta,self.data_phi_kk,self.data_f_kk,self.data_phif2_kk)

        m_phif2_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_f2_mass,kk_f2_width,kk_f2_const,kk_f2_theta,self.mc_phi_kk,self.mc_f_kk,self.mc_phif2_kk)
        
        d_u_rho_r_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_u_rho_mass,kk_u_rho_width,kk_u_rho_const,kk_u_rho_theta,self.data_phi_kk,self.data_b123_kk,self.data_u_rho_r_kk)

        m_u_rho_r_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_u_rho_mass,kk_u_rho_width,kk_u_rho_const,kk_u_rho_theta,self.mc_phi_kk,self.mc_b123_kk,self.mc_u_rho_r_kk)
        
        d_u_rho_l_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_u_rho_mass,kk_u_rho_width,kk_u_rho_const,kk_u_rho_theta,self.data_phi_kk,self.data_b124_kk,self.data_u_rho_l_kk)

        m_u_rho_l_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_u_rho_mass,kk_u_rho_width,kk_u_rho_const,kk_u_rho_theta,self.mc_phi_kk,self.mc_b124_kk,self.mc_u_rho_l_kk)
        
        d_u_b_r_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_u_b_mass,kk_u_b_width,kk_u_b_const,kk_u_b_theta,self.data_phi_kk,self.data_b124_kk,self.data_u_b_r_kk)

        m_u_b_r_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_u_b_mass,kk_u_b_width,kk_u_b_const,kk_u_b_theta,self.mc_phi_kk,self.mc_b124_kk,self.mc_u_b_r_kk)
        
        d_u_b_l_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_u_b_mass,kk_u_b_width,kk_u_b_const,kk_u_b_theta,self.data_phi_kk,self.data_b123_kk,self.data_u_b_l_kk)

        m_u_b_l_kk_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,kk_u_b_mass,kk_u_b_width,kk_u_b_const,kk_u_b_theta,self.mc_phi_kk,self.mc_b123_kk,self.mc_u_b_l_kk)
        
        d_phif0_pipi_BW_flatte980 = self.calculate_BW_flatte980(phi_mass,phi_width,pipi_f980_mass,pipi_f980_g_pipi,pipi_f980_rg,pipi_f980_const,pipi_f980_theta,self.data_phi_pipi,self.data_f_pipi,self.data_phif0_pipi)

        m_phif0_pipi_BW_flatte980 = self.calculate_BW_flatte980(phi_mass,phi_width,pipi_f980_mass,pipi_f980_g_pipi,pipi_f980_rg,pipi_f980_const,pipi_f980_theta,self.mc_phi_pipi,self.mc_f_pipi,self.mc_phif0_pipi)
        
        d_phif0_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_f0_mass,pipi_f0_width,pipi_f0_const,pipi_f0_theta,self.data_phi_pipi,self.data_f_pipi,self.data_phif0_pipi)

        m_phif0_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_f0_mass,pipi_f0_width,pipi_f0_const,pipi_f0_theta,self.mc_phi_pipi,self.mc_f_pipi,self.mc_phif0_pipi)
        
        d_phif2_pipi_BW_flatte1270 = self.calculate_BW_flatte1270(phi_mass,phi_width,pipi_f1270_mass,pipi_f1270_width,pipi_f1270_const,pipi_f1270_theta,self.data_phi_pipi,self.data_f_pipi,self.data_phif2_pipi)

        m_phif2_pipi_BW_flatte1270 = self.calculate_BW_flatte1270(phi_mass,phi_width,pipi_f1270_mass,pipi_f1270_width,pipi_f1270_const,pipi_f1270_theta,self.mc_phi_pipi,self.mc_f_pipi,self.mc_phif2_pipi)
        
        d_phif2_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_f2_mass,pipi_f2_width,pipi_f2_const,pipi_f2_theta,self.data_phi_pipi,self.data_f_pipi,self.data_phif2_pipi)

        m_phif2_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_f2_mass,pipi_f2_width,pipi_f2_const,pipi_f2_theta,self.mc_phi_pipi,self.mc_f_pipi,self.mc_phif2_pipi)
        
        d_u_rho_r_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_u_rho_mass,pipi_u_rho_width,pipi_u_rho_const,pipi_u_rho_theta,self.data_phi_pipi,self.data_b123_pipi,self.data_u_rho_r_pipi)

        m_u_rho_r_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_u_rho_mass,pipi_u_rho_width,pipi_u_rho_const,pipi_u_rho_theta,self.mc_phi_pipi,self.mc_b123_pipi,self.mc_u_rho_r_pipi)
        
        d_u_rho_l_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_u_rho_mass,pipi_u_rho_width,pipi_u_rho_const,pipi_u_rho_theta,self.data_phi_pipi,self.data_b124_pipi,self.data_u_rho_l_pipi)

        m_u_rho_l_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_u_rho_mass,pipi_u_rho_width,pipi_u_rho_const,pipi_u_rho_theta,self.mc_phi_pipi,self.mc_b124_pipi,self.mc_u_rho_l_pipi)
        
        d_u_b_r_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_u_b_mass,pipi_u_b_width,pipi_u_b_const,pipi_u_b_theta,self.data_phi_pipi,self.data_b124_pipi,self.data_u_b_r_pipi)

        m_u_b_r_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_u_b_mass,pipi_u_b_width,pipi_u_b_const,pipi_u_b_theta,self.mc_phi_pipi,self.mc_b124_pipi,self.mc_u_b_r_pipi)
        
        d_u_b_l_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_u_b_mass,pipi_u_b_width,pipi_u_b_const,pipi_u_b_theta,self.data_phi_pipi,self.data_b123_pipi,self.data_u_b_l_pipi)

        m_u_b_l_pipi_BW_BW = self.calculate_BW_BW(phi_mass,phi_width,pipi_u_b_mass,pipi_u_b_width,pipi_u_b_const,pipi_u_b_theta,self.mc_phi_pipi,self.mc_b123_pipi,self.mc_u_b_l_pipi)