BUG: Bambi logistic regression fails on M1 silicon

[tarmojuristo]

Describe the issue:

Running a logistic regression in Bambi with default settings fails. However, if chains=1 is set in model.fit() then everything works fine.

Reproduceable code example:

import bambi as bmb
import numpy as np

data = bmb.load_data("adults")

age_mean = np.mean(data["age"])
age_std = np.std(data["age"])

data["age"] = (data["age"] - age_mean) / age_std

model = bmb.Model("income['>50K'] ~ sex + age", data, family="bernoulli")
fitted = model.fit(draws=1000)

Error message:

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [Intercept, sex, age]
Traceback (most recent call last):000 00:00<? Sampling 4 chains, 0 divergences]
  File "/Users/tarmo/Desktop/SALK/streamlit-model/example.py", line 12, in <module>
    fitted = model.fit(draws=1000)
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/site-packages/bambi/models.py", line 271, in fit
    return self.backend.run(
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/site-packages/bambi/backend/pymc.py", line 98, in run
    result = self._run_mcmc(
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/site-packages/bambi/backend/pymc.py", line 289, in _run_mcmc
    idata = pm.sample(
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/site-packages/pymc/sampling/mcmc.py", line 529, in sample
    mtrace = _mp_sample(**sample_args, **parallel_args)
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/site-packages/pymc/sampling/mcmc.py", line 1019, in _mp_sample
    for draw in sampler:
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/site-packages/pymc/sampling/parallel.py", line 449, in __iter__
    draw = ProcessAdapter.recv_draw(self._active)
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/site-packages/pymc/sampling/parallel.py", line 320, in recv_draw
    msg = ready[0].recv()
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/multiprocessing/connection.py", line 255, in recv
    buf = self._recv_bytes()
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/multiprocessing/connection.py", line 419, in _recv_bytes
    buf = self._recv(4)
  File "/Users/tarmo/mambaforge/envs/pymc/lib/python3.10/multiprocessing/connection.py", line 388, in _recv
    raise EOFError
EOFError

PyMC version information:

PyMC 5.0.1
PyTensor 2.8.11
Python 3.10
MacOS Ventura 13.0.1 on M1
mambaforge

Context for the issue:

No response

[twiecki]

Does a super simple model

with pm.Model():
    x = pm.Normal("x")
    pm.sample()

also produce that error?

[tarmojuristo]

no, this works fine. also the bambi t-test example samples without problems.

[tomicapretto]

@tarmojuristo could you try with this example?

import numpy as np
import pymc as pm

size = 100
rng = np.random.default_rng(1234)
x = rng.normal(size=size)
g = rng.choice(list("ABC"), size=size)
y = rng.integers(0, 1, size=size, endpoint=True)

g_levels, g_idxs = np.unique(g, return_inverse=True)
coords = {"group": g_levels}

with pm.Model(coords=coords) as model:
    coef_x = pm.Normal("x")
    coef_g = pm.Normal("g", dims="group")
    p = pm.math.softmax(coef_x + coef_g[g_idxs])
    pm.Bernoulli("y", p=p, observed=y)
    idata = pm.sample()

[tarmojuristo]

all good: Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.

[tomicapretto]

If you agree, I can give you a couple of examples you can try until we find the broken part.

What if you do

import bambi as bmb
import numpy as np
import pymc as pm

data = bmb.load_data("adults")

age_mean = np.mean(data["age"])
age_std = np.std(data["age"])

data["age"] = (data["age"] - age_mean) / age_std

model = bmb.Model("income['>50K'] ~ sex + age", data, family="bernoulli")
model.build()
with model.backend.model:
    idata = pm.sample()

[tarmojuristo]

this one fails with the same error

[tomicapretto]

Ok! Later today I'll give another chunk of PyMC code to test. I'll try to reproduce the PyMC model created by Bambi more closely

[michaelosthege]

Make sure to avoid model code on the 0th indent level. This can lead to recursion problems with multiprocessing depending on the OS-dependent fork/spawn setting.

Just wrap everything in a simple run() function:

def run():
    # do stuff


if __name__ == "__main__":
    run()

[tarmojuristo]

Make sure to avoid model code on the 0th indent level. This can lead to recursion problems with multiprocessing depending on the OS-dependent fork/spawn setting.

Just wrap everything in a simple run() function:

def run():
    # do stuff


if __name__ == "__main__":
    run()

Thanks, this is a useful tip - definitely avoids some trouble on my intel-based macbook - but on M1 the above example still fails, even when wrapped into a function.

[ivanmkc]

+1, I get this error as well when using a Jupyter notebook.

[tomicapretto]

import bambi as bmb
import pandas as pd
import numpy as np
import pymc as pm
import pytensor.tensor as pt

from formulae import design_matrices

data = bmb.load_data("adults")

age_mean = np.mean(data["age"])
age_std = np.std(data["age"])

data["age"] = (data["age"] - age_mean) / age_std

def logit(x):
    """Logit function that ensures result is in (0, 1)"""
    eps = np.finfo(float).eps
    result = pt.sigmoid(x)
    result = pt.switch(pt.eq(result, 0), eps, result)
    result = pt.switch(pt.eq(result, 1), 1 - eps, result)
    return result

dm = design_matrices("income['>50K'] ~ sex + age", data)
y = np.array(dm.response)
X = np.array(dm.common)

coords = {"sex_dim": ["Male"]}
with pm.Model(coords=coords) as pm_model:
    eta = 0
    intercept = pm.Normal("intercept", mu=0, sigma=4.34)
    sex = pt.atleast_1d(pm.Normal("sex", mu=0, sigma=5.31, dims="sex_dim"))
    age = pt.atleast_1d(pm.Normal("age", mu=0, sigma=2.5, shape=None))
    
    # Concatenate parameters
    coefs = pt.concatenate([sex, age])
    
    # Create linear predictor
    eta += intercept 
    eta += pt.dot(X[:, 1:], coefs) # drop first column because that's for the intercept

    # Use inverse of link function
    p = logit(eta)

    pm.Bernoulli("income", p=p, observed=y)

This resembles much more closely what's happening under the hood. Could you run it and let me know if it works or fails?

pm.model_to_graphviz(pm_model)

[tarmojuristo]

When I sample this one, I once again get the EOFError.

[tomicapretto]

Now I think we can try "removing" parts that may be conflicting. I would start with the dims for sex or the replacing the custom logit function with pt.sigmoid