Heyo, I am new to JAX and I am trying to make my code jit-compatible so it runs faster on big data arrays. Here are the functions I am trying to write using JAX: @jax.jit def logL(p, clustering): same = clustering[:, None] == clustering[None, :] m = (p*same + (1-p)*~same) return (m.sum() - jnp.diag(m).sum())/2 @partial(jax.jit, static_argnames="n_timepoints") def LD(clusterings, candidate, scm_neuron, n_timepoints): def body(i, delta): conf_mat = scm_neuron[:, i] new = logL(conf_mat, candidate[:, i]) old = logL(conf_mat, clusterings[:, i]) return delta + (new - old) deltas= jax.lax.fori_loop(0, int(n_timepoints), body, 0.0) return deltas These two functions are parts of MCMC algorithm, where I sample certain clustering labels for neurons. For each of these clusterings, I calculate the logL in accordance with a predefined probability matrix (p). The acceptance step (whether to move to the next sample, or stay at the current one) depends on the delta which is calculated by the second function LD. Here, all the input arrays have a shape (n_neurons, n_timepoints). I use jax.lax.fori_loop to go over every timepoint and accumulate the deltas in the final output. However, this seems to be slower than a python for-loop. Profiler (Profiler output) shows that nearly half of the time is dedicated to slicing operation, as far as I understand. I do not understand why it takes so much time for slicing, since my jnp.arrays are relatively small. I also tried to rewrite it using jax.lax.scan and jax.vmap, however, it does not solve it. I wonder how I can overcome this. Thanks a lot in advance! Continue reading...
