search

pymc-devs / pymc

Bayesian Modeling and Probabilistic Programming in Python
https://docs.pymc.io/
Other
8.68k stars 2.01k forks source link

BUG: CPython crashes within sampling (code 141) #7126

Open YishuangW opened 8 months ago

YishuangW commented 8 months ago

Describe the issue:

When running the following code, CPython crashed with exit code 141. There won't be any issue if I only use one chain to sample. It only occurs when I use more than one chain.

Reproduceable code example:

# Import some useful modules.
import jax.numpy as jnp
import os

# Import JAX-FEM specific modules.
from jax_fem.problem import Problem
from jax_fem.solver import solver
from jax_fem.utils import save_sol
from jax_fem.generate_mesh import box_mesh, get_meshio_cell_type, Mesh
from jax_fem import logger
import logging
logger.setLevel(logging.ERROR)

import pymc as pm
import numpy as np
import arviz as az
import torch
import matplotlib.pyplot as plt
import pytensor.tensor as pt

def main():
# Define constitutive relationship.
    class LinearElasticity(Problem):
        # The function 'get_tensor_map' overrides base class method. Generally, JAX-FEM
        # solves -div(f(u_grad)) = b. Here, we have f(u_grad) = sigma.
        theta: int = 0
        def get_tensor_map(self):
            def stress(u_grad):
                E = self.theta
                nu = 0.3
                mu = E / (2. * (1. + nu))
                lmbda = E * nu / ((1 + nu) * (1 - 2 * nu))
                epsilon = 0.5 * (u_grad + u_grad.T)
                sigma = lmbda * jnp.trace(epsilon) * jnp.eye(
                    self.dim) + 2 * mu * epsilon
                return sigma
            return stress

        def get_surface_maps(self):
            def surface_map(u, x):
                return np.array([0., 0., -100.])
            return [surface_map]

    # Specify mesh-related information (second-order tetrahedron element).
    ele_type = 'TET10'
    cell_type = get_meshio_cell_type(ele_type)
    data_dir = os.path.join(os.path.dirname(__file__), 'data')
    Lx, Ly, Lz = 10., 2., 2.
    Nx, Ny, Nz = 5, 1, 1
    meshio_mesh = box_mesh(Nx=Nx,
                           Ny=Ny,
                           Nz=Nz,
                           Lx=Lx,
                           Ly=Ly,
                           Lz=Lz,
                           data_dir=data_dir,
                           ele_type=ele_type)
    mesh = Mesh(meshio_mesh.points, meshio_mesh.cells_dict[cell_type])

    # Define boundary locations.
    def left(point):
        return jnp.isclose(point[0], 0., atol=1e-5)

    def right(point):
        return jnp.isclose(point[0], Lx, atol=1e-5)

    # Define Dirichlet boundary values.
    # This means on the 'left' side, we apply the function 'zero_dirichlet_val'
    # to all components of the displacement variable u.
    def zero_dirichlet_val(point):
        return 0.

    dirichlet_bc_info = [[left] * 3, [0, 1, 2], [zero_dirichlet_val] * 3]

    # Define Neumann boundary locations.
    # This means on the 'right' side, we will perform the surface integral to get
    # the tractions with the function 'get_surface_maps' defined in the class 'LinearElasticity'.
    location_fns = [right]

    # Create an instance of the problem.
    problem = LinearElasticity(mesh,
                               vec=3,
                               dim=3,
                               ele_type=ele_type,
                               dirichlet_bc_info=dirichlet_bc_info,
                               location_fns=location_fns)
    # Solve the defined problem.
    problem.theta = 70e3
    sol = solver(problem, linear=True, use_petsc=True)

    def my_model(theta):
        theta = np.asarray(theta)
        theta = theta.item()
        problem.theta = theta
        sol = solver(problem, linear=True, use_petsc=True)
        sol = np.array(sol)
        # sol = torch.from_numpy(sol)
        return sol

    # JAX_TRACEBACK_FILTERING=off

    def my_loglike(theta, data, sigma):
        model = my_model(theta)
        # n = theta.shape[0]
        return -(0.5 / sigma**2) * np.sum((data - model) ** 2)

    class LogLike(pt.Op):

        """
        Specify what type of object will be passed and returned to the Op when it is
        called. In our case we will be passing it a vector of values (the parameters
        that define our model) and returning a single "scalar" value (the
        log-likelihood)
        """

        itypes = [pt.dvector]  # expects a vector of parameter values when called
        otypes = [pt.dscalar]  # outputs a single scalar value (the log likelihood)

        def __init__(self, loglike, data, sigma):
            """
            Initialise the Op with various things that our log-likelihood function
            requires. Below are the things that are needed in this particular
            example.

            Parameters
            ----------
            loglike:
                The log-likelihood (or whatever) function we've defined
            data:
                The "observed" data that our log-likelihood function takes in
            x:
                The dependent variable (aka 'x') that our model requires
            sigma:
                The noise standard deviation that our function requires.
            """

            # add inputs as class attributes
            self.likelihood = loglike
            self.data = data
            self.sigma = sigma

        def perform(self, node, inputs, outputs):
            # the method that is used when calling the Op
            theta = inputs  # this will contain my variables

            # call the log-likelihood function
            logl = self.likelihood(theta, self.data, self.sigma)

            outputs[0][0] = np.array(logl)  # output the log-likelihood

    # set up our data
    sigma = 0.05  # standard deviation of noise
    data = sol
    # create our Op
    logl = LogLike(my_loglike, data, sigma)

    # use PyMC to sampler from log-likelihood

    with pm.Model():
        # uniform priors on m and c
        param1 = pm.Normal("param1", mu=60000, sigma=1)
        # TODO: POISSON RATIO
        # TODO:: multivariate normal

        # convert m and c to a tensor vector
        theta = pt.as_tensor_variable([param1])

        # use a Potential to "call" the Op and include it in the logp computation
        pm.Potential("likelihood", logl(theta))

        # Use custom number of draws to replace the HMC based defaults
        # Won't crash only if chain is equal to 1
        idata_mh = pm.sample(10, tune=5, chains=3, mp_ctx="forkserver") # CRASH HERE!

        print(os.getcwd(), flush=True)
        print("\n\n\n\n\nFinished sampling\n\n\n\n\n\n\n\n", flush=True)
        torch.save(idata_mh, 'file_gp_real.pt')
        print("\n\n\n\n\nFinished saving torch\n\n\n\n\n\n\n\n", flush=True)
        az.plot_trace(idata_mh, compact=False, combined=False)
        print("\n\n\n\n\nFinished plotting trace\n\n\n\n\n\n\n", flush=True)
        plt.savefig("result_gp_real.png")
        print("\n\n\n\n\nFinished saving fig\n\n\n\n\n\n\n\n", flush=True)

if __name__ == '__main__':
    __spec__ = None
    main()
    exit(0)

Error message:

No error message from python since it's crashed. Here's the stack trace captured by gdb:

#0  __GI___libc_write (nbytes=26, buf=0x7f545f11c450, fd=15) at ../sysdeps/unix/sysv/linux/write.c:26
#1  __GI___libc_write (fd=fd@entry=15, buf=buf@entry=0x7f545f11c450, nbytes=nbytes@entry=26) at ../sysdeps/unix/sysv/linux/write.c:24
#2  0x00005598acfb94ac in _Py_write_impl (gil_held=1, count=26, buf=0x7f545f11c450, fd=15) at /usr/local/src/conda/python-3.9.18/Python/fileutils.c:1771
#3  _Py_write (fd=15, buf=0x7f545f11c450, count=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/fileutils.c:1831
#4  0x00005598acfe0773 in os_write_impl (data=0x7ffc88d73160, fd=15, module=<optimized out>) at /usr/local/src/conda/python-3.9.18/Modules/posixmodule.c:9768
#5  os_write (module=<optimized out>, args=0x7f5410d52a98, nargs=<optimized out>) at /usr/local/src/conda/python-3.9.18/Modules/clinic/posixmodule.c.h:5257
#6  0x00005598acef2e73 in cfunction_vectorcall_FASTCALL (func=0x7f56b0eadbd0, args=0x7f5410d52a98, nargsf=<optimized out>, kwnames=<optimized out>)
    at /usr/local/src/conda/python-3.9.18/Objects/methodobject.c:430
#7  0x00005598acee2232 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d52a98, callable=0x7f56b0eadbd0, tstate=0x5598ad3e7cf0)
    at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#8  PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d52a98, callable=0x7f56b0eadbd0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#9  call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#10 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f5410d52900, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3520
#11 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x7f5410d52900, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#12 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, 
    kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x7f563e3fda30, qualname=0x7f562c4e4170)
    at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#13 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>)
    at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#14 0x00005598acee24e6 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d57cb8, callable=0x7f562c4efca0, tstate=0x5598ad3e7cf0)
    at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#15 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d57cb8, callable=0x7f562c4efca0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#16 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#17 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f5410d57b20, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3506
#18 0x00005598acef2712 in _PyEval_EvalFrame (throwflag=0, f=0x7f5410d57b20, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#19 function_code_fastcall (tstate=0x5598ad3e7cf0, co=0x7f562c4e3920, args=<optimized out>, nargs=2, globals=0x7f562c4d8d00) at /usr/local/src/conda/python-3.9.18/Objects/call.c:330
#20 0x00005598acee24e6 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d52700, callable=0x7f562c4efdc0, tstate=0x5598ad3e7cf0)
    at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#21 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d52700, callable=0x7f562c4efdc0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#22 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#23 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f5410d52580, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3506
#24 0x00005598acef2712 in _PyEval_EvalFrame (throwflag=0, f=0x7f5410d52580, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#25 function_code_fastcall (tstate=0x5598ad3e7cf0, co=0x7f562c4dfb30, args=<optimized out>, nargs=2, globals=0x7f562c4d8d00) at /usr/local/src/conda/python-3.9.18/Objects/call.c:330
#26 0x00005598acee24e6 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x5598bac68400, callable=0x7f562c4ef820, tstate=0x5598ad3e7cf0)
    at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#27 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x5598bac68400, callable=0x7f562c4ef820) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#28 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#29 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x5598bac68260, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3506
#30 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x5598bac68260, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#31 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, 
    kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x7f563e3fda30, qualname=0x7f563e401120)
    at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#32 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>)
    at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#33 0x00005598acee24e6 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d58eb8, callable=0x7f563e41d3a0, tstate=0x5598ad3e7cf0)
    at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#34 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d58eb8, callable=0x7f563e41d3a0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#35 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#36 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f5410d58d40, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3506
#37 0x00005598acef2712 in _PyEval_EvalFrame (throwflag=0, f=0x7f5410d58d40, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#38 function_code_fastcall (tstate=0x5598ad3e7cf0, co=0x7f563e3febe0, args=<optimized out>, nargs=1, globals=0x7f563e3fc9c0) at /usr/local/src/conda/python-3.9.18/Objects/call.c:330
#39 0x00005598acee24e6 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d599a0, callable=0x7f563e41d550, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#40 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d599a0, callable=0x7f563e41d550) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#41 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#42 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f5410d59800, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3506
#43 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x7f5410d59800, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#44 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x7f563e3fdef0, qualname=0x7f563e4014e0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#45 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#46 0x00005598acee690e in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d5d378, callable=0x7f563e41d790, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#47 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f5410d5d378, callable=0x7f563e41d790) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#48 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#49 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f5410d5d1f0, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3489
#50 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x7f5410d5d1f0, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#51 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x7f56b0f308f0, qualname=0x7f563e4017b0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#52 0x00005598aceff8cb in _PyFunction_Vectorcall (kwnames=<optimized out>, nargsf=<optimized out>, stack=0x5598baa4d368, func=0x7f563e41ddc0) at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#53 _PyObject_VectorcallTstate (kwnames=<optimized out>, nargsf=<optimized out>, args=0x5598baa4d368, callable=0x7f563e41ddc0, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#54 method_vectorcall (method=<optimized out>, args=0x5598baa4d370, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/classobject.c:53
#55 0x00005598acee2232 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x5598baa4d370, callable=0x7f542c995fc0, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#56 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x5598baa4d370, callable=0x7f542c995fc0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#57 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#58 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x5598baa4d170, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3520
#59 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x5598baa4d170, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#60 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x7f563e426180, closure=0x0, name=0x7f563e3f76b0, qualname=0x7f563e3f76b0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#61 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#62 0x00005598acf0004c in PyVectorcall_Call (kwargs=<optimized out>, tuple=<optimized out>, callable=0x7f563e421dc0) at /usr/local/src/conda/python-3.9.18/Objects/call.c:243
#63 _PyObject_Call (kwargs=<optimized out>, args=<optimized out>, callable=0x7f563e421dc0, tstate=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:266
#64 PyObject_Call (callable=0x7f563e421dc0, args=<optimized out>, kwargs=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:293
#65 0x00005598acee63f6 in do_call_core (kwdict=0x7f5410dc4080, callargs=0x7f56b0f13040, func=0x7f563e421dc0, tstate=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5125
#66 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x5598b8bbf200, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3582
#67 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x5598b8bbf200, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#68 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x7f563e4255c0, closure=0x0, name=0x7f56b03b39f0, qualname=0x7f56b03b39f0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#69 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#70 0x00005598acee3042 in _PyObject_VectorcallTstate (kwnames=0x7f56b029e440, nargsf=<optimized out>, args=<optimized out>, callable=0x7f563e421a60, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#71 PyObject_Vectorcall (kwnames=0x7f56b029e440, nargsf=<optimized out>, args=<optimized out>, callable=0x7f563e421a60) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#72 call_function (kwnames=0x7f56b029e440, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#73 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x5598b488e600, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3537
#74 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x5598b488e600, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#75 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x7f56b0e8d470, qualname=0x7f56b0e8d470) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#76 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#77 0x00005598acee2232 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f56b051d690, callable=0x7f56b0380790, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#78 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f56b051d690, callable=0x7f56b0380790) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#79 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#80 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f56b051d520, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3520
#81 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x7f56b051d520, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#82 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=2, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x0, qualname=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#83 0x00005598acee0c88 in _PyEval_EvalCodeWithName (_co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=0x0, kwcount=0, kwstep=2, defs=0x0, defcount=0, kwdefs=0x0, closure=0x0, name=0x0, qualname=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4361
#84 0x00005598acee0c39 in PyEval_EvalCodeEx (_co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kws=<optimized out>, kwcount=0, defs=0x0, defcount=0, kwdefs=0x0, closure=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4377
#85 0x00005598acf8e2cb in PyEval_EvalCode (co=co@entry=0x7f56b029f920, globals=globals@entry=0x7f56b0e7ff40, locals=locals@entry=0x7f56b0e7ff40) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:828
#86 0x00005598acf92dd2 in builtin_exec_impl (module=<optimized out>, locals=0x7f56b0e7ff40, globals=0x7f56b0e7ff40, source=0x7f56b029f920) at /usr/local/src/conda/python-3.9.18/Python/bltinmodule.c:1026
#87 builtin_exec (module=<optimized out>, args=<optimized out>, nargs=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/clinic/bltinmodule.c.h:396
#88 0x00005598acef2e73 in cfunction_vectorcall_FASTCALL (func=0x7f56b0ee2040, args=0x7f56b037da20, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/methodobject.c:430
#89 0x00005598acf8e352 in _PyObject_VectorcallTstate (tstate=0x5598ad3e7cf0, callable=0x7f56b0ee2040, args=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#90 0x00005598ace1c6e1 in call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5074
#91 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f56b037d8b0, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3520
#92 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x7f56b037d8b0, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#93 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=2, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x0, qualname=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#94 0x00005598acee0c88 in _PyEval_EvalCodeWithName (_co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=0x0, kwcount=0, kwstep=2, defs=0x0, defcount=0, kwdefs=0x0, closure=0x0, name=0x0, qualname=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4361
#95 0x00005598acee0c39 in PyEval_EvalCodeEx (_co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kws=<optimized out>, kwcount=0, defs=0x0, defcount=0, kwdefs=0x0, closure=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4377
#96 0x00005598acf8e2cb in PyEval_EvalCode (co=co@entry=0x7f56b03a53a0, globals=globals@entry=0x7f56b0e7ff40, locals=locals@entry=0x7f56b0e7ff40) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:828
#97 0x00005598acf92dd2 in builtin_exec_impl (module=<optimized out>, locals=0x7f56b0e7ff40, globals=0x7f56b0e7ff40, source=0x7f56b03a53a0) at /usr/local/src/conda/python-3.9.18/Python/bltinmodule.c:1026
#98 builtin_exec (module=<optimized out>, args=<optimized out>, nargs=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/clinic/bltinmodule.c.h:396
#99 0x00005598acef2e73 in cfunction_vectorcall_FASTCALL (func=0x7f56b0ee2040, args=0x7f56b04fc3c8, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/methodobject.c:430
#100 0x00005598acf8e352 in _PyObject_VectorcallTstate (tstate=0x5598ad3e7cf0, callable=0x7f56b0ee2040, args=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#101 0x00005598ace1c6e1 in call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5074
#102 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f56b04fc230, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3520
#103 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x7f56b04fc230, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#104 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x7f56b0ed4b30, qualname=0x7f56b04f4470) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#105 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#106 0x00005598acee24e6 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f56b04fcd78, callable=0x7f56b03d5b80, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#107 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f56b04fcd78, callable=0x7f56b03d5b80) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#108 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#109 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f56b04fcbe0, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3506
#110 0x00005598acef2712 in _PyEval_EvalFrame (throwflag=0, f=0x7f56b04fcbe0, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#111 function_code_fastcall (tstate=0x5598ad3e7cf0, co=0x7f56b0392920, args=<optimized out>, nargs=2, globals=0x7f56b037b040) at /usr/local/src/conda/python-3.9.18/Objects/call.c:330
#112 0x00005598acee24e6 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x5598ad503428, callable=0x7f56b0398700, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#113 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x5598ad503428, callable=0x7f56b0398700) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#114 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#115 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x5598ad503260, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3506
#116 0x00005598acef2712 in _PyEval_EvalFrame (throwflag=0, f=0x5598ad503260, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#117 function_code_fastcall (tstate=0x5598ad3e7cf0, co=0x7f56b03940e0, args=<optimized out>, nargs=0, globals=0x7f56b037b040) at /usr/local/src/conda/python-3.9.18/Objects/call.c:330
#118 0x00005598acee690e in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x7f56b0b83a78, callable=0x7f56b0398c10, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#119 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x7f56b0b83a78, callable=0x7f56b0398c10) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#120 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#121 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x7f56b0b83900, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3489
#122 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x7f56b0b83900, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#123 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=2, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x0, qualname=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#124 0x00005598acee0c88 in _PyEval_EvalCodeWithName (_co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=0x0, kwcount=0, kwstep=2, defs=0x0, defcount=0, kwdefs=0x0, closure=0x0, name=0x0, qualname=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4361
#125 0x00005598acee0c39 in PyEval_EvalCodeEx (_co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kws=<optimized out>, kwcount=0, defs=0x0, defcount=0, kwdefs=0x0, closure=0x0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4377
#126 0x00005598acf8e2cb in PyEval_EvalCode (co=co@entry=0x7f56b04dfdf0, globals=globals@entry=0x7f56b0e7ff40, locals=locals@entry=0x7f56b0e7ff40) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:828
#127 0x00005598acf92dd2 in builtin_exec_impl (module=<optimized out>, locals=0x7f56b0e7ff40, globals=0x7f56b0e7ff40, source=0x7f56b04dfdf0) at /usr/local/src/conda/python-3.9.18/Python/bltinmodule.c:1026
#128 builtin_exec (module=<optimized out>, args=<optimized out>, nargs=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/clinic/bltinmodule.c.h:396
#129 0x00005598acef2e73 in cfunction_vectorcall_FASTCALL (func=0x7f56b0ee2040, args=0x5598ad4ae6d0, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/methodobject.c:430
#130 0x00005598acee2232 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x5598ad4ae6d0, callable=0x7f56b0ee2040, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#131 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x5598ad4ae6d0, callable=0x7f56b0ee2040) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#132 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#133 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x5598ad4ae510, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3520
#134 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x5598ad4ae510, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#135 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x7f56b0b9afb0, qualname=0x7f56b0b9afb0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#136 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#137 0x00005598acee2232 in _PyObject_VectorcallTstate (kwnames=0x0, nargsf=<optimized out>, args=0x5598ad486b18, callable=0x7f56b0b9b9d0, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:118
#138 PyObject_Vectorcall (kwnames=0x0, nargsf=<optimized out>, args=0x5598ad486b18, callable=0x7f56b0b9b9d0) at /usr/local/src/conda/python-3.9.18/Include/cpython/abstract.h:127
#139 call_function (kwnames=0x0, oparg=<optimized out>, pp_stack=<synthetic pointer>, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:5077
#140 _PyEval_EvalFrameDefault (tstate=<optimized out>, f=0x5598ad486970, throwflag=<optimized out>) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:3520
#141 0x00005598acee0fb4 in _PyEval_EvalFrame (throwflag=0, f=0x5598ad486970, tstate=0x5598ad3e7cf0) at /usr/local/src/conda/python-3.9.18/Include/internal/pycore_ceval.h:40
#142 _PyEval_EvalCode (tstate=<optimized out>, _co=<optimized out>, globals=<optimized out>, locals=<optimized out>, args=<optimized out>, argcount=<optimized out>, kwnames=<optimized out>, kwargs=<optimized out>, kwcount=<optimized out>, kwstep=1, defs=<optimized out>, defcount=<optimized out>, kwdefs=0x0, closure=0x0, name=0x7f56b0b9d2b0, qualname=0x7f56b0b9d2b0) at /usr/local/src/conda/python-3.9.18/Python/ceval.c:4329
#143 0x00005598acef2469 in _PyFunction_Vectorcall (func=<optimized out>, stack=<optimized out>, nargsf=<optimized out>, kwnames=<optimized out>) at /usr/local/src/conda/python-3.9.18/Objects/call.c:396
#144 0x00005598acfaf298 in pymain_run_module (modname=<optimized out>, set_argv0=1) at /usr/local/src/conda/python-3.9.18/Modules/main.c:295
#145 0x00005598acfae80d in pymain_run_python (exitcode=0x7ffc88d75e30) at /usr/local/src/conda/python-3.9.18/Modules/main.c:596
#146 Py_RunMain () at /usr/local/src/conda/python-3.9.18/Modules/main.c:681
#147 0x00005598acf82107 in Py_BytesMain (argc=<optimized out>, argv=<optimized out>) at /usr/local/src/conda/python-3.9.18/Modules/main.c:1101
#148 0x00007f56b0c29d90 in __libc_start_call_main (main=main@entry=0x5598acf820c0 <main>, argc=argc@entry=4, argv=argv@entry=0x7ffc88d76068) at ../sysdeps/nptl/libc_start_call_main.h:58
#149 0x00007f56b0c29e40 in __libc_start_main_impl (main=0x5598acf820c0 <main>, argc=4, argv=0x7ffc88d76068, init=<optimized out>, fini=<optimized out>, rtld_fini=<optimized out>, stack_end=0x7ffc88d76058) at ../csu/libc-start.c:392
#150 0x00005598acf82001 in _start ()


### PyMC version information:

PyMC version v5.10.3
Build hd8ed1ab_0
Installed from conda-forge 
Ubantu 22.04.3 LTS.

### Context for the issue:

I am unable to do parallel sampling with PyMC, which greatly slows my work.
welcome[bot] commented 8 months ago

Welcome Banner] :tada: Welcome to PyMC! :tada: We're really excited to have your input into the project! :sparkling_heart:
If you haven't done so already, please make sure you check out our Contributing Guidelines and Code of Conduct.

ricardoV94 commented 8 months ago

For JAX code you may have better luck using a full-jax based sampler like numpyro. We describe this in https://www.pymc.io/projects/examples/en/latest/howto/wrapping_jax_function.html#unwrapping-the-wrapped-jax-function