Assigning multiple CPU cores for some functions

[samarth-kashyap]

Hi jax team,

I'm attempting an eigenvalue optimization problem using numpyro. I have a matrix whose eigenvalues depend on 3 parameters (w1, w3, w5). The eig_func_mcmc constructs the matrix and solves the eigenvalue problem. Numpyro uses MCMC and builds multiple chains which randomly sample the parameters (w1, w3, w5).

def model():
    # setting min and max value to be 0.1*wt and 3.*wt
    w1min, w1max = .1*abs(w1t), 3.*abs(w1t)
    w3min, w3max = .1*abs(w3t), 3.*abs(w3t)
    w5min, w5max = .1*abs(w5t), 3.*abs(w5t)
    
    w1 = numpyro.sample('w1', dist.Uniform(w1min, w1max))
    w3 = numpyro.sample('w3', dist.Uniform(w3min, w3max))
    w5 = numpyro.sample('w5', dist.Uniform(w5min, w5max))
    
    sigma = numpyro.sample('sigma', dist.Uniform(0.1, 10.0))
    eig_sample = numpyro.deterministic('eig', eig_mcmc_func(w1=w1, w3=w3, w5=w5))
    return numpyro.sample('obs', dist.Normal(eig_sample, sigma), obs=eigvals_true)

Based on some tests, the eigenvalue problem solving on the CPU speeds up significantly when the number of cores is increased to 8 and marginally thereafter. I would want to setup the code such that each individual eigenvalue problem solving would take 8 cores. Is there a way to do this?

Numpyro handles parallelizing of chains using pmap. So ideally what I want to do is to assign 8*N cores to a job, with N chains
such that numpyro executes each chain on a different block of 8 cores. Thanks!