Use jit_compile with numpy in Pyro

I have a question on using the “jit_compile = True” in NUTS algorithm with a numpy model.
My numpy model calculates chlorine 36 concentration given a seismic scenario attributed by pyro. Basically I detach the tensors containing the samples and cast them in numpy arrays.Then I turn the numpy array containing the resulting chlorine 36 concentration into a torch tensor.
For now, my code works fine but is very slow (it requires days to weeks to complete). I noticed that using the jit argument speed up the inferrence but returns errors and wrong results. Is there some commands to tune the jit_compile to work that way?
I hope I am being clear, thanks in advance.

i’m sorry but i don’t follow. you can’t arbitrarily mix pyro and numpy and pytorch and detach and expect things to work so you’ll need to be much more specific.

Ok here’s the main code :

import forward_function as forward
import geometric_scaling_factors
from constants import constants
import torch
import pyro
import numpy as np
import pyro.distributions as dist
import post_process as fig
from seismic_scenario import seismic_scenario as true_scenario
from pyro.infer import MCMC, NUTS
import parameters

""" Input seismic scenario """
erosion_rate = 0 # Erosion rate (mm/yr)
number_of_events = 3
seismic_scenario['erosion_rate'] = erosion_rate

""" Input parameters"""
cl36AMS = param.cl36AMS
height = param.h
Hfinal = param.Hfinal
Data = torch.tensor(cl36AMS)

""" Geometric scaling """
scaling_depth_rock, scaling_depth_coll, scaling_surf_rock, scaling_factors = geometric_scaling_factors.neutron_scaling(param, constants, number_of_events+1)

""" MCMC parameters """
w_step = 50 # number of warmup
nb_sample = 100 # number of samples

""" MCMC model """
def model(obs):
    ages = torch.zeros((number_of_events))
    ages[0] = pyro.sample('age1', dist.Uniform(2.0, 30*1e3))

    for i in range (1, number_of_events):
        max_age = ages[i-1]
        ages[i] = pyro.sample('age'+str(i+1), dist.Uniform(2.0, max_age))  
    # print('\n age', ages)
    slips = torch.zeros((number_of_events))
    slips[0]=pyro.sample('slip1', dist.Uniform(0.0, Hfinal-100))
    for i in range (1, number_of_events-1):
        slips[i]=pyro.sample('slip'+str(i+1), dist.Uniform(0.0, max_slip))

    seismic_scenario['ages'] = ages
    seismic_scenario['slips'] = slips
    seismic_scenario['SR'] = true_scenario['SR']
    seismic_scenario['preexp'] = true_scenario['preexp']
    seismic_scenario['quiescence'] = true_scenario['quiescence']
    sigma=pyro.sample('sigma', dist.Uniform(0, 10000)) 

    t = forward.mds_torch(seismic_scenario, scaling_factors, constants, parameters, 500, 200)
    return pyro.sample('obs', dist.Normal(t, sigma), obs=obs)

""" usage MCMC """
kernel = NUTS(model) # chose kernel (NUTS, HMC, ...)
mcmc = MCMC(kernel, warmup_steps=w_step, num_samples=nb_sample)
posterior_samples = mcmc.get_samples()
print('MCMC done \n')

Inside the mds_torch module I detach the tensors to perform calculations with numpy (which is much faster) and then return a torch tensor. My question is : is it possible to do the same thing with jit_compile=True in the kernel (maybe with a different method)?

I tested this with synthetic data without the jit_compile, it works but slowly, and I am looking for ways to speed up the process.
I have tried to work with torch in the mds_torch module, but it is even slower and as of today, never reached the end of first iteration.

I hope it is clearer, thanks again.