Repository: pymc-devs/pymc
Pull request: Use jaxified logp for initial point evaluation when sampling via Jax
Use jaxified logp for initial point evaluation when sampling via Jax

## Description
- get jaxified logp function in `sample_jax_nuts`
  - uses different parameters to get jaxified function depending on which `nuts_sampler` is specified
- pass jaxified logp function to `_get_batched_jittered_initial_points`
  - added `logp_fn` parameter to function signature
  - wrap passed function to conform to how `_init_jitter` will call it
- pass wrapped function to `_init_jitter`
  - added `logp_fn` parameter to function signature
- added logic in `_init_jitter` to decide which function to use to evaluate the generated points
- added a bunch of type annotations

## Related Issue
- [x] Closes #7608


## Checklist
- [x] Checked that [the pre-commit linting/style checks pass](https://docs.pymc.io/en/latest/contributing/python_style.html)
- [x] Included tests that prove the fix is effective or that the new feature works
- [x] Added necessary documentation (docstrings and/or example notebooks) - have a question about the `sample_blackjax_nuts` function docstring, will put in comment below
- [x] If you are a pro: each commit corresponds to a [relevant logical change]

## Type of change
- [x] New feature / enhancement
- [x] Bug fix
- [ ] Documentation
- [x] Maintenance
- [ ] Other (please specify):

<!-- readthedocs-preview pymc start -->
----
📚 Documentation preview 📚: https://pymc--7610.org.readthedocs.build/en/7610/

<!-- readthedocs-preview pymc end -->

Problem statement:
BUG: model initial_point fails when `pt.config.floatX = "float32"`
### Describe the issue:

I have a hierarchical mmm model setup (with pymc, not pymc-marketing) and have been successfully using it with `float32`s up to the 5.19 release. It is using numpyro/jax to sample.

With 5.19 I am getting errors in `_init_jitter`, it appears that there is something going wrong when passing the generated initial points to the compiled logp function. I think the use of zerosumNormal distributions is causing the problem but I'm not sure if it's the values returned by the `ipfn(seed)` or the evaluation in the compiled `model_logp_fn`. I've included the verbose model debug return from my model, but the reproducible example below is using the example radon model.

### Reproduceable code example:

```python
import pytensor as pt

pt.config.floatX = "float32"
pt.config.warn_float64 = "ignore"

import numpy as np

import pandas as pd
import pymc as pm

# Load the radon dataset
data = pd.read_csv(pm.get_data("radon.csv"))
data["log_radon"] = data["log_radon"].astype(np.float64)
county_idx, counties = pd.factorize(data.county)
coords = {"county": counties, "obs_id": np.arange(len(county_idx))}

# Create a simple hierarchical model for the radon dataset
with pm.Model(coords=coords, check_bounds=False) as model:
    intercept = pm.Normal("intercept", sigma=10)

    # County effects
    raw = pm.ZeroSumNormal("county_raw", dims="county")
    sd = pm.HalfNormal("county_sd")
    county_effect = pm.Deterministic("county_effect", raw * sd, dims="county")

    # Global floor effect
    floor_effect = pm.Normal("floor_effect", sigma=2)

    # County:floor interaction
    raw = pm.ZeroSumNormal("county_floor_raw", dims="county")
    sd = pm.HalfNormal("county_floor_sd")
    county_floor_effect = pm.Deterministic("county_floor_effect", raw * sd, dims="county")

    mu = (
        intercept
        + county_effect[county_idx]
        + floor_effect * data.floor.values
        + county_floor_effect[county_idx] * data.floor.values
    )

    sigma = pm.HalfNormal("sigma", sigma=1.5)
    pm.Normal("log_radon", mu=mu, sigma=sigma, observed=data.log_radon.values, dims="obs_id")

idata = pm.sample(
    model=model,
    chains=1,
    tune=500,
    draws=500,
    progressbar=False,
    compute_convergence_checks=False,
    return_inferencedata=False,
    # compile_kwargs=dict(mode="NUMBA")
)
```


### Error message:

```shell
SamplingError: Initial evaluation of model at starting point failed!
Starting values:
{'mu_adstock_logodds__': array([-0.05812943, -1.3623805 , -0.11238763, -0.5489    , -1.0595012 ,
       -1.6993128 , -1.0827137 , -1.0727112 , -1.7406836 , -0.25875893,
       -1.0392814 , -0.94655335,  0.08540058, -0.5202998 ,  0.33476347,
       -0.5190295 , -1.0631133 , -0.57824194, -0.7651843 , -0.87049246],
      dtype=float32), 'mu_lambda_log__': array([0.80796826, 0.6380648 , 1.2703241 , 1.9139447 , 1.5351636 ,
       1.7267554 , 1.0714266 , 1.0567162 , 2.0499995 , 1.5639257 ,
       1.9248102 , 1.6098787 , 0.56195414, 0.22962454, 0.10311574,
       0.36890286, 0.88685906, 1.7213213 , 1.2186754 , 0.36423793],
      dtype=float32), 'mu_a': array(0.37686896, dtype=float32), 'z_a_state_zerosum__': array([ 0.19204606,  0.04281855,  0.5793297 , -0.2614823 , -0.60941106,
       -0.5645286 ], dtype=float32), 'z_a_age_zerosum__': array([-0.39018095,  0.82710207,  0.93642265, -0.30185702,  0.7038302 ],
      dtype=float32), 'z_a_brand_zerosum__': array([-0.19112755], dtype=float32), 'z_a_cohort_zerosum__': array([0.48081324], dtype=float32), 'roas_rv_log__': array([-2.0530562 , -1.1738657 , -2.1438377 , -0.5091131 , -1.7475982 ,
       -0.71832424,  0.7603461 , -0.6206605 , -0.06742464, -1.7866528 ,
       -0.8273122 ,  0.3449509 , -0.4612695 , -1.5151714 , -0.9103267 ,
        0.92829466,  0.05409501, -0.17818904,  0.4767271 ,  1.0659611 ],
      dtype=float32), 'z_b_state_zerosum__': array([[ 0.9404972 ,  0.5173597 ,  0.947904  ,  0.13234697,  0.91516215,
         0.8013973 ],
       [ 0.76825845, -0.00395296,  0.20068735,  0.89948386, -0.8211762 ,
        -0.85832894],
       [-0.23005143, -0.21612945,  0.38762948, -0.17819382, -0.12716204,
         0.5923743 ],
       [ 0.42042527, -0.8524518 ,  0.41847345, -0.27271616,  0.9246246 ,
         0.10817041],
       [ 0.5263319 ,  0.2471719 , -0.7485617 , -0.30133858,  0.48334756,
         0.3941832 ],
       [ 0.20216979,  0.85309595, -0.13996926,  0.13978173,  0.2398329 ,
        -0.6436887 ],
       [ 0.6737369 ,  0.56847805, -0.6305385 ,  0.11638433, -0.88574356,
        -0.28407305],
       [-0.88218284, -0.5086114 ,  0.5899615 ,  0.4378485 ,  0.1587185 ,
        -0.06953551],
       [-0.4915884 ,  0.5323545 ,  0.5476521 ,  0.03720003,  0.93754584,
         0.20691545],
       [ 0.52630943,  0.8224987 ,  0.38608572, -0.48465434, -0.43573558,
         0.8957741 ],
       [-0.36157402, -0.80625093, -0.35036987,  0.58878934, -0.01390497,
        -0.9381669 ],
       [ 0.84330225, -0.24921258, -0.5905653 ,  0.71469283,  0.04288541,
        -0.15973687],
       [ 0.16558522, -0.82091933,  0.6832055 ,  0.9772107 ,  0.05962521,
        -0.7585791 ],
       [-0.34562954, -0.14827304, -0.7131723 , -0.49016312,  0.92288506,
         0.5976615 ],
       [ 0.14006369, -0.78234476,  0.58030427,  0.7087098 ,  0.83111507,
        -0.4035646 ],
       [-0.89755225, -0.7397433 ,  0.87145287,  0.5502826 , -0.75381684,
        -0.08390825],
       [ 0.52683276,  0.2059658 , -0.65518147,  0.9072355 ,  0.02874351,
         0.8622515 ],
       [-0.31716648, -0.22920816, -0.4773452 ,  0.24883471, -0.71415335,
        -0.00118549],
       [ 0.67181724,  0.83293927,  0.06912381,  0.59113485,  0.15970528,
         0.15737481],
       [ 0.91586804, -0.00439743,  0.8568587 , -0.5599965 ,  0.9024629 ,
        -0.92843384]], dtype=float32), 'z_b_age_zerosum__': array([[ 2.76717365e-01, -5.32886267e-01, -4.37204897e-01,
         6.65345013e-01,  9.51409161e-01],
       [ 5.68430305e-01,  5.80501974e-01,  9.02473092e-01,
        -1.05089314e-01, -6.80017248e-02],
       [-7.40278482e-01, -6.54472530e-01,  5.73029280e-01,
        -2.49546096e-01, -9.83492434e-01],
       [ 4.39474136e-01,  4.55991507e-01,  2.91431248e-01,
         7.93459237e-01,  8.33085358e-01],
       [ 7.54141450e-01, -9.88980412e-01,  7.37549663e-01,
        -9.54164326e-01,  3.69425505e-01],
       [-4.45214868e-01,  1.18648775e-01,  4.35143918e-01,
         7.96567798e-01, -5.37025869e-01],
       [-8.19233358e-01, -3.28816384e-01, -4.24525887e-01,
        -4.72912073e-01, -2.51088679e-01],
       [ 4.16932464e-01,  1.86953232e-01, -2.34448180e-01,
        -5.28278828e-01, -7.83707380e-01],
       [-2.08375111e-01, -1.69877082e-01, -9.20472383e-01,
        -5.55105388e-01,  2.24135935e-01],
       [-7.18319640e-02, -8.23212624e-01, -1.14380375e-01,
         3.75080615e-01,  4.38587993e-01],
       [-7.84464240e-01,  9.64653268e-02,  1.33498237e-01,
         2.63148099e-01,  9.03292537e-01],
       [ 3.70964020e-01,  8.96655202e-01, -9.78391707e-01,
         6.00353360e-01, -3.29210430e-01],
       [ 9.80111659e-01,  6.26725018e-01,  8.71558905e-01,
        -7.08010912e-01,  3.21216695e-02],
       [-8.55567873e-01, -4.15038317e-01, -2.70858496e-01,
         7.64281690e-01,  1.69419169e-01],
       [-6.30245388e-01,  5.22969842e-01, -6.22790098e-01,
         8.40588808e-01,  5.42818129e-01],
       [-2.61249971e-02,  5.77672958e-01, -9.52823997e-01,
        -5.49517214e-01, -4.92883384e-01],
       [ 3.19638193e-01,  8.80902350e-01,  2.54505854e-02,
        -4.16665673e-01,  7.45047331e-01],
       [-1.37079775e-01, -9.72663925e-04,  2.88793862e-01,
         9.96275783e-01,  5.16300082e-01],
       [-2.26764768e-01, -9.21454072e-01, -2.66458213e-01,
        -2.89255470e-01, -5.44357836e-01],
       [-7.08415627e-01, -2.39693552e-01,  1.69611976e-01,
         6.88308775e-01, -2.90724158e-01]], dtype=float32), 'z_b_brand_zerosum__': array([[ 0.5283236 ],
       [-0.14984424],
       [ 0.5653758 ],
       [ 0.6604869 ],
       [ 0.5594151 ],
       [ 0.36363953],
       [-0.13847719],
       [ 0.2760732 ],
       [ 0.60931265],
       [ 0.39675766],
       [ 0.13061056],
       [-0.843226  ],
       [-0.24025336],
       [ 0.21590135],
       [ 0.38261482],
       [-0.9853659 ],
       [-0.89518636],
       [-0.73512644],
       [ 0.24093248],
       [ 0.53579485]], dtype=float32), 'z_b_cohort_zerosum__': array([[-0.89432555],
       [ 0.56194794],
       [ 0.8784441 ],
       [ 0.29107055],
       [-0.03464561],
       [-0.7969992 ],
       [ 0.3803715 ],
       [ 0.14672193],
       [-0.72870535],
       [-0.9154726 ],
       [-0.05452913],
       [ 0.3046088 ],
       [ 0.11125816],
       [ 0.47531176],
       [-0.26566276],
       [-0.97336334],
       [ 0.8083869 ],
       [ 0.10919274],
       [ 0.34626842],
       [-0.14968246]], dtype=float32), 'mu_b_pos_con': array([-0.55350554, -1.271216  , -2.4323664 , -2.2574682 , -1.7525793 ],
      dtype=float32), 'z_b_pos_con_state_zerosum__': array([[ 0.7088774 , -0.64973927,  0.01091878,  0.42494458, -0.5738007 ,
        -0.13388953],
       [ 0.2020803 , -0.4415184 ,  0.9188914 , -0.91647035,  0.20404132,
        -0.19665638],
       [-0.460112  ,  0.250589  ,  0.43331107, -0.31884784, -0.20407897,
        -0.85575414],
       [ 0.053673  , -0.31756902,  0.743054  , -0.83611   ,  0.24574716,
         0.31937018],
       [-0.24382843, -0.7048149 , -0.29253975, -0.2640002 , -0.7875447 ,
         0.9551751 ]], dtype=float32), 'z_b_pos_con_age_zerosum__': array([[-0.968712  , -0.18256   ,  0.54134816,  0.5418771 , -0.20875828],
       [ 0.07226439,  0.3872751 ,  0.78559935,  0.56196845, -0.21816622],
       [-0.07638574,  0.56413966, -0.29065445, -0.8537547 ,  0.26876295],
       [-0.03170992,  0.7177836 , -0.14550653,  0.9839035 ,  0.05719684],
       [-0.09502391,  0.7952581 , -0.3214874 ,  0.5325462 , -0.84900486]],
      dtype=float32), 'z_b_pos_con_brand_zerosum__': array([[ 0.9145228 ],
       [-0.576082  ],
       [-0.9886196 ],
       [-0.76018703],
       [ 0.8999498 ]], dtype=float32), 'z_b_pos_con_cohort_zerosum__': array([[-0.32495368],
       [-0.32046455],
       [ 0.20123298],
       [-0.7780437 ],
       [-0.78692883]], dtype=float32), 'mu_b_neg_con': array([-1.6263036], dtype=float32), 'z_b_neg_con_state_zerosum__': array([[ 0.6405368 ,  0.02604678,  0.7836906 , -0.9528276 ,  0.65443355,
        -0.00763485]], dtype=float32), 'z_b_neg_con_age_zerosum__': array([[-0.88560593, -0.3676332 , -0.1298519 , -0.39139104,  0.05216115]],
      dtype=float32), 'z_b_neg_con_brand_zerosum__': array([[0.06288974]], dtype=float32), 'z_b_neg_con_cohort_zerosum__': array([[-0.2893259]], dtype=float32), 'mu_b_lag': array([-3.4769938], dtype=float32), 'z_b_lag_state_zerosum__': array([[-0.5544768 ,  0.8843655 , -0.85776746, -0.54040647,  0.83116657,
        -0.91366553]], dtype=float32), 'z_b_lag_age_zerosum__': array([[-0.27069455,  0.2308255 ,  0.02965806, -0.922125  ,  0.25014895]],
      dtype=float32), 'z_b_lag_brand_zerosum__': array([[-0.95589054]], dtype=float32), 'z_b_lag_cohort_zerosum__': array([[-0.9150893]], dtype=float32), 'mu_b_fourier_year': array([ 0.99425113,  0.6344655 , -0.11898322,  0.8841637 ,  0.21039897,
        0.07808983,  0.8532348 ,  0.44000137,  0.98853546, -0.52862114,
        0.79177517, -0.01937965, -0.7798197 , -0.8342347 , -0.64850366,
        0.97407717,  0.95000124,  0.7092928 ,  0.6113028 , -0.27019796],
      dtype=float32), 'sd_y_log__': array(4.4781275, dtype=float32)}

Logp initial evaluation results:
{'mu_adstock': -20.54, 'mu_lambda': -213.9, 'mu_a': -1.3, 'z_a_state': -inf, 'z_a_age': -inf, 'z_a_brand': -0.44, 'z_a_cohort': -inf, 'roas_rv': -72.16, 'z_b_state': -inf, 'z_b_age': -inf, 'z_b_brand': -inf, 'z_b_cohort': -inf, 'mu_b_pos_con': -5.82, 'z_b_pos_con_state': -inf, 'z_b_pos_con_age': -inf, 'z_b_pos_con_brand': -inf, 'z_b_pos_con_cohort': -66.75, 'mu_b_neg_con': -1.12, 'z_b_neg_con_state': -inf, 'z_b_neg_con_age': -inf, 'z_b_neg_con_brand': 1.19, 'z_b_neg_con_cohort': -2.8, 'mu_b_lag': -1.03, 'z_b_lag_state': -inf, 'z_b_lag_age': -inf, 'z_b_lag_brand': -44.3, 'z_b_lag_cohort': -40.49, 'mu_b_fourier_year': -91.0, 'sd_y': -1.21, 'y_like': -56160.32}
You can call `model.debug()` for more details.

SamplingError                             Traceback (most recent call last)
File ~/.ipykernel/1917/command--1-2479103431:18
     15 entry = [ep for ep in metadata.distribution("mmm_v2").entry_points if ep.name == "mmm"]
     16 if entry:
     17   # Load and execute the entrypoint, assumes no parameters
---> 18   entry[0].load()()
     19 else:
     20   import importlib

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/mmm_v2/main.py:16, in main()
     13 pt.config.floatX = "float32"  # pyright: ignore[reportPrivateImportUsage] // TODO: we can probably set this via env vars or pytensor.rc
     14 config = setup_mmm()
---> 16 run_mmm(config)

-- elided non important dataprep pipeline stuff --

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/mmm_v2/model/ModelBuilder.py:634, in HLM_ModelBuilder.fit(self)
    631 """Fits the model with the pymc model `sample()` method."""
    632 self.logger.debug("Sampling model.")
    633 self.model_trace.extend(
--> 634     pm.sample(
    635         draws=self.config.draws,
    636         tune=self.config.tune,
    637         chains=self.config.chains,
    638         model=self.model,
    639         nuts_sampler="numpyro",
    640         idata_kwargs={"log_likelihood": False},
    641         var_names=[f"{x}" for x in self.model.free_RVs],
    642         target_accept=self.config.target_accept,
    643         random_seed=self.config.sampler_seed,
    644         progressbar=self.config.progress_bars,
    645     )
    646 )
    647 self.logger.debug("Finished sampling.")
    649 HLM_ModelBuilder._print_diagnostics(self.model_trace, self.logger)

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/pymc/sampling/mcmc.py:773, in sample(draws, tune, chains, cores, random_seed, progressbar, progressbar_theme, step, var_names, nuts_sampler, initvals, init, jitter_max_retries, n_init, trace, discard_tuned_samples, compute_convergence_checks, keep_warning_stat, return_inferencedata, idata_kwargs, nuts_sampler_kwargs, callback, mp_ctx, blas_cores, model, compile_kwargs, **kwargs)
    768         raise ValueError(
    769             "Model can not be sampled with NUTS alone. It either has discrete variables or a non-differentiable log-probability."
    770         )
    772     with joined_blas_limiter():
--> 773         return _sample_external_nuts(
    774             sampler=nuts_sampler,
    775             draws=draws,
    776             tune=tune,
    777             chains=chains,
    778             target_accept=kwargs.pop("nuts", {}).get("target_accept", 0.8),
    779             random_seed=random_seed,
    780             initvals=initvals,
    781             model=model,
    782             var_names=var_names,
    783             progressbar=progressbar,
    784             idata_kwargs=idata_kwargs,
    785             compute_convergence_checks=compute_convergence_checks,
    786             nuts_sampler_kwargs=nuts_sampler_kwargs,
    787             **kwargs,
    788         )
    790 if exclusive_nuts and not provided_steps:
    791     # Special path for NUTS initialization
    792     if "nuts" in kwargs:

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/pymc/sampling/mcmc.py:389, in _sample_external_nuts(sampler, draws, tune, chains, target_accept, random_seed, initvals, model, var_names, progressbar, idata_kwargs, compute_convergence_checks, nuts_sampler_kwargs, **kwargs)
    386 elif sampler in ("numpyro", "blackjax"):
    387     import pymc.sampling.jax as pymc_jax
--> 389     idata = pymc_jax.sample_jax_nuts(
    390         draws=draws,
    391         tune=tune,
    392         chains=chains,
    393         target_accept=target_accept,
    394         random_seed=random_seed,
    395         initvals=initvals,
    396         model=model,
    397         var_names=var_names,
    398         progressbar=progressbar,
    399         nuts_sampler=sampler,
    400         idata_kwargs=idata_kwargs,
    401         compute_convergence_checks=compute_convergence_checks,
    402         **nuts_sampler_kwargs,
    403     )
    404     return idata
    406 else:

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/pymc/sampling/jax.py:595, in sample_jax_nuts(draws, tune, chains, target_accept, random_seed, initvals, jitter, model, var_names, nuts_kwargs, progressbar, keep_untransformed, chain_method, postprocessing_backend, postprocessing_vectorize, postprocessing_chunks, idata_kwargs, compute_convergence_checks, nuts_sampler)
    589 vars_to_sample = list(
    590     get_default_varnames(filtered_var_names, include_transformed=keep_untransformed)
    591 )
    593 (random_seed,) = _get_seeds_per_chain(random_seed, 1)
--> 595 initial_points = _get_batched_jittered_initial_points(
    596     model=model,
    597     chains=chains,
    598     initvals=initvals,
    599     random_seed=random_seed,
    600     jitter=jitter,
    601 )
    603 if nuts_sampler == "numpyro":
    604     sampler_fn = _sample_numpyro_nuts

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/pymc/sampling/jax.py:225, in _get_batched_jittered_initial_points(model, chains, initvals, random_seed, jitter, jitter_max_retries)
    209 def _get_batched_jittered_initial_points(
    210     model: Model,
    211     chains: int,
   (...)
    215     jitter_max_retries: int = 10,
    216 ) -> np.ndarray | list[np.ndarray]:
    217     """Get jittered initial point in format expected by NumPyro MCMC kernel.
    218 
    219     Returns
   (...)
    223         Each item has shape `(chains, *var.shape)`
    224     """
--> 225     initial_points = _init_jitter(
    226         model,
    227         initvals,
    228         seeds=_get_seeds_per_chain(random_seed, chains),
    229         jitter=jitter,
    230         jitter_max_retries=jitter_max_retries,
    231     )
    232     initial_points_values = [list(initial_point.values()) for initial_point in initial_points]
    233     if chains == 1:

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/pymc/sampling/mcmc.py:1382, in _init_jitter(model, initvals, seeds, jitter, jitter_max_retries, logp_dlogp_func)
   1379 if not np.isfinite(point_logp):
   1380     if i == jitter_max_retries:
   1381         # Print informative message on last attempted point
-> 1382         model.check_start_vals(point)
   1383     # Retry with a new seed
   1384     seed = rng.integers(2**30, dtype=np.int64)

File /local_disk0/.ephemeral_nfs/cluster_libraries/python/lib/python3.10/site-packages/pymc/model/core.py:1769, in Model.check_start_vals(self, start, **kwargs)
   1766 initial_eval = self.point_logps(point=elem, **kwargs)
   1768 if not all(np.isfinite(v) for v in initial_eval.values()):
-> 1769     raise SamplingError(
   1770         "Initial evaluation of model at starting point failed!\n"
   1771         f"Starting values:\n{elem}\n\n"
   1772         f"Logp initial evaluation results:\n{initial_eval}\n"
   1773         "You can call `model.debug()` for more details."
   1774     )
```


### PyMC version information:

Running on a windows machine in a linux container
pymc installed via poetry(pypi).
using `libopenblas`

annotated-types            0.7.0
arviz                      0.19.0
babel                      2.16.0
blinker                    1.4
build                      1.2.2.post1
CacheControl               0.14.1
cachetools                 5.5.0
certifi                    2024.8.30
cfgv                       3.4.0
charset-normalizer         3.4.0
cleo                       2.1.0
click                      8.1.7
cloudpickle                3.1.0
colorama                   0.4.6
cons                       0.4.6
contourpy                  1.3.1
coverage                   7.6.8
crashtest                  0.4.1
cryptography               3.4.8
cycler                     0.12.1
dbus-python                1.2.18
distlib                    0.3.9
distro                     1.7.0
distro-info                1.1+ubuntu0.2
dm-tree                    0.1.8
dulwich                    0.21.7
etuples                    0.3.9
exceptiongroup             1.2.2
fastjsonschema             2.21.1
filelock                   3.16.1
fonttools                  4.55.1
ghp-import                 2.1.0
graphviz                   0.20.3
griffe                     1.5.1
h5netcdf                   1.4.1
h5py                       3.12.1
httplib2                   0.20.2
identify                   2.6.3
idna                       3.10
importlib_metadata         8.5.0
iniconfig                  2.0.0
installer                  0.7.0
jaraco.classes             3.4.0
jax                        0.4.35
jaxlib                     0.4.35
jeepney                    0.7.1
Jinja2                     3.1.4
joblib                     1.4.2
keyring                    24.3.1
kiwisolver                 1.4.7
launchpadlib               1.10.16
lazr.restfulclient         0.14.4
lazr.uri                   1.0.6
logical-unification        0.4.6
loguru                     0.7.2
Markdown                   3.7
markdown-it-py             3.0.0
MarkupSafe                 3.0.2
matplotlib                 3.9.3
mdurl                      0.1.2
mergedeep                  1.3.4
miniKanren                 1.0.3
mkdocs                     1.6.1
mkdocs-autorefs            1.2.0
mkdocs-gen-files           0.5.0
mkdocs-get-deps            0.2.0
mkdocs-glightbox           0.4.0
mkdocs-literate-nav        0.6.1
mkdocs-material            9.5.47
mkdocs-material-extensions 1.3.1
mkdocs-section-index       0.3.9
mkdocstrings               0.26.2
mkdocstrings-python        1.12.2
ml_dtypes                  0.5.0
mmm_v2                     0.0.1         /workspaces/mmm_v2
more-itertools             8.10.0
msgpack                    1.1.0
multimethod                1.10
multipledispatch           1.0.0
mypy-extensions            1.0.0
nodeenv                    1.9.1
numpy                      1.26.4
numpyro                    0.15.3
oauthlib                   3.2.0
opt_einsum                 3.4.0
packaging                  24.2
paginate                   0.5.7
pandas                     2.2.3
pandera                    0.20.4
pathspec                   0.12.1
pexpect                    4.9.0
pillow                     11.0.0
pip                        24.3.1
pkginfo                    1.12.0
platformdirs               4.3.6
pluggy                     1.5.0
poetry                     1.8.4
poetry-core                1.9.1
poetry-plugin-export       1.8.0
pre_commit                 4.0.1
ptyprocess                 0.7.0
pyarrow                    18.1.0
pydantic                   2.10.3
pydantic_core              2.27.1
Pygments                   2.18.0
PyGObject                  3.42.1
PyJWT                      2.3.0
pymc                       5.19.1
pymc-marketing             0.6.0
pymdown-extensions         10.12
pyparsing                  3.2.0
pyproject_hooks            1.2.0
pytensor                   2.26.4
pytest                     8.3.4
pytest-cov                 6.0.0
python-apt                 2.4.0+ubuntu3
python-dateutil            2.9.0.post0
pytz                       2024.2
PyYAML                     6.0.2
pyyaml_env_tag             0.1
RapidFuzz                  3.10.1
regex                      2024.11.6
requests                   2.32.3
requests-toolbelt          1.0.0
rich                       13.9.4
ruff                       0.8.1
scikit-learn               1.5.2
scipy                      1.14.1
seaborn                    0.13.2
SecretStorage              3.3.1
setuptools                 75.6.0
shellingham                1.5.4
six                        1.17.0
ssh-import-id              5.11
threadpoolctl              3.5.0
tomli                      2.2.1
tomlkit                    0.13.2
toolz                      1.0.0
tqdm                       4.67.1
trove-classifiers          2024.10.21.16
typeguard                  4.4.1
typing_extensions          4.12.2
typing-inspect             0.9.0
tzdata                     2024.2
unattended-upgrades        0.1
urllib3                    2.2.3
virtualenv                 20.28.0
wadllib                    1.3.6
watchdog                   6.0.0
wheel                      0.38.4
wrapt                      1.17.0
xarray                     2024.9.0
xarray-einstats            0.8.0
zipp                       3.21.0

### Context for the issue:

I want to use f32 due to memory constraints, I've got a big model and dataset with a hierarchy that takes up 40+gb of ram when running using f32 so I'd need a huge box to go up to f64.

Maybe the optimisations in 5.19 make that problem moot? I've managed to get it running with f64 temporarily but I'm not sure if it'll be a long term solution

Hints from issue discussion:
![Welcome Banner](https://raw.githubusercontent.com/pymc-devs/brand/main/welcome-bot/BannerWelcome.jpg)]
:tada: Welcome to _PyMC_! :tada: We're really excited to have your input into the project! :sparkling_heart:
<br>If you haven't done so already, please make sure you check out our [Contributing Guidelines](https://www.pymc.io/projects/docs/en/latest/contributing/index.html) and [Code of Conduct](https://github.com/pymc-devs/pymc/blob/main/CODE_OF_CONDUCT.md).

debug trace wouldn't fit in the issue:

``` shell
Model debug info from using the debugger to pause at mcmc.py L1386 and calling `model.debug(point, verbose=True)`

point={'mu_adstock_logodds__': array([-0.4352023 , -0.2916657 , -0.22959015, -0.13479875, -0.30916345,
       -1.5283642 , -1.655654  , -0.48334104, -0.30794948,  0.33021173,
        0.0656831 , -1.6596988 , -0.47080922, -0.91784286, -0.15092045,
       -1.1535047 , -1.0793369 , -0.8457366 , -1.3219562 , -1.3434141 ],
      dtype=float32), 'mu_lambda_log__': array([2.2058203 , 1.453856  , 1.6157596 , 0.43113965, 1.6577175 ,
       0.84230906, 0.27920187, 0.899368  , 1.8292453 , 2.2519255 ,
       2.3814824 , 2.108992  , 1.4626849 , 1.5980046 , 0.8299414 ,
       0.3984745 , 1.8697526 , 1.4488894 , 2.1219566 , 0.18294996],
      dtype=float32), 'mu_a': array(0.16189706, dtype=float32), 'z_a_state_zerosum__': array([-0.28566697, -0.48986882, -0.24404879, -0.4398191 , -0.85556257,
        0.10748474], dtype=float32), 'z_a_age_zerosum__': array([-0.8873828 , -0.32871515,  0.57389575,  0.30648488, -0.6199051 ],
      dtype=float32), 'z_a_brand_zerosum__': array([-0.46795005], dtype=float32), 'z_a_cohort_zerosum__': array([0.4478183], dtype=float32), 'roas_rv_log__': array([-0.6555429 , -0.068955  , -0.9309296 , -0.47975117, -1.3743148 ,
        0.47711676,  0.14864558,  0.10893539,  0.38675046, -0.8370424 ,
       -0.66264176, -1.5664616 , -0.572411  , -1.4761899 , -1.0794916 ,
       -0.2966979 ,  0.05331168,  1.0579    ,  0.13020533,  0.8946314 ],
      dtype=float32), 'z_b_state_zerosum__': array([[-0.00928065,  0.22661608,  0.26506114,  0.10868097,  0.47448924,
         0.6638588 ],
       [ 0.39453208,  0.9125831 , -0.05330055,  0.16203904, -0.7428469 ,
        -0.2878859 ],
       [-0.22169496,  0.25026786,  0.7721265 , -0.16524933, -0.8161399 ,
         0.5124463 ],
       [ 0.4635996 ,  0.935813  ,  0.3664374 , -0.39854062, -0.11831979,
         0.23826346],
       [ 0.4144873 , -0.07588482, -0.4675184 ,  0.9954296 ,  0.44995347,
         0.6562674 ],
       [ 0.94758034,  0.2068893 ,  0.6966277 ,  0.31964955, -0.8013234 ,
        -0.59591883],
       [-0.17715056, -0.7038275 ,  0.18067661,  0.01431344, -0.9491178 ,
        -0.3321023 ],
       [ 0.3942007 ,  0.9996393 , -0.31270924, -0.08990093, -0.09300919,
        -0.16450764],
       [-0.00447197, -0.61609423,  0.8628801 ,  0.96006954,  0.7203218 ,
        -0.7518324 ],
       [ 0.6942957 , -0.44699988,  0.57910615,  0.8879041 ,  0.531556  ,
        -0.9510816 ],
       [ 0.78471005,  0.10752742,  0.4335172 , -0.58196217, -0.9966123 ,
        -0.17337854],
       [-0.06716231, -0.5351729 , -0.1103561 , -0.15798165, -0.15524508,
         0.8739795 ],
       [ 0.47066316,  0.03429028, -0.2272006 ,  0.57281727,  0.9989922 ,
        -0.26203355],
       [-0.59414744, -0.34866652, -0.58397436, -0.12034182,  0.16198853,
        -0.36454397],
       [ 0.12944746, -0.05762197,  0.99427617,  0.81767935,  0.5921547 ,
         0.9800794 ],
       [ 0.9717736 ,  0.9814946 ,  0.4856121 , -0.5534532 ,  0.11700594,
         0.9247631 ],
       [-0.2042932 , -0.411241  ,  0.27332023,  0.9046378 ,  0.6154953 ,
        -0.08056752],
       [ 0.9214974 , -0.65947914, -0.41038954,  0.54713   , -0.3560202 ,
        -0.9969325 ],
       [-0.08087815,  0.18727091,  0.84307253, -0.48801887,  0.29456693,
         0.5796735 ],
       [ 0.27902   ,  0.29730567, -0.40406513, -0.18478568, -0.61452186,
        -0.5549851 ]], dtype=float32), 'z_b_age_zerosum__': array([[-6.8486583e-01, -4.6815168e-02, -5.9378707e-01,  8.9602423e-01,
         8.5502052e-01],
       [ 7.8728044e-01, -3.6670679e-01, -3.3962426e-01,  3.0838227e-01,
         3.6406529e-01],
       [ 2.6335692e-02, -6.4281446e-01,  5.1187193e-01,  8.4743094e-01,
         2.2725777e-01],
       [ 1.9795492e-01,  9.2090023e-01, -8.8563585e-01, -2.8022802e-01,
        -2.3639840e-01],
       [ 2.9900335e-02,  7.1486712e-01,  6.7400551e-01, -7.0308822e-01,
        -6.9614536e-01],
       [ 1.3013636e-01,  1.0248652e-01, -5.7761997e-02, -8.4077924e-01,
        -4.7718164e-01],
       [-5.9512955e-01, -9.2812777e-02, -9.9525869e-02,  9.1666229e-02,
         7.0176089e-01],
       [ 4.7530079e-01,  2.4512438e-01, -4.2329890e-01, -6.4359361e-01,
         4.8717073e-01],
       [-4.3700787e-01, -3.6620468e-01, -5.4181212e-01, -7.7344787e-01,
        -5.0397778e-01],
       [ 5.0274485e-01,  8.3750290e-01, -1.0284081e-01,  6.3057953e-01,
         9.5303126e-02],
       [ 9.8040715e-02,  3.1213897e-01,  9.7941196e-01,  6.9000393e-01,
         1.3390434e-01],
       [ 5.9936064e-01, -1.6784413e-01,  9.4419844e-02, -5.2747607e-02,
         1.4664505e-01],
       [-8.8135488e-02, -1.6365480e-01, -2.4431337e-01,  2.9717276e-01,
         8.4692138e-01],
       [ 4.3495792e-01,  5.3633377e-02,  4.0893257e-01, -1.4952035e-01,
        -2.0427135e-01],
       [-4.2102399e-01,  7.0554173e-01, -7.1471161e-01, -6.7319351e-01,
         6.4274400e-01],
       [ 2.9349172e-01, -4.3267983e-01, -6.5261596e-01, -7.2232783e-01,
         8.7439209e-01],
       [ 3.5815242e-01,  9.3956250e-01,  7.0418483e-01,  6.0373771e-01,
         3.7868690e-02],
       [-8.4775686e-01, -4.1210197e-02, -8.3802587e-01, -7.5553125e-01,
         1.8493308e-01],
       [-9.3824327e-01,  4.9592870e-01,  7.3176724e-01,  1.4875463e-01,
         4.6959123e-01],
       [-9.1654237e-04, -3.3683771e-01,  2.9216877e-01, -1.3542533e-01,
         3.8034424e-02]], dtype=float32), 'z_b_brand_zerosum__': array([[ 0.1755699 ],
       [-0.64566314],
       [-0.164115  ],
       [-0.7532449 ],
       [ 0.8172519 ],
       [-0.07537084],
       [ 0.18247645],
       [-0.6951736 ],
       [-0.34786522],
       [ 0.13041314],
       [-0.26256517],
       [-0.75953436],
       [ 0.13290128],
       [ 0.7542543 ],
       [ 0.43538788],
       [ 0.2650157 ],
       [-0.22280337],
       [ 0.9436325 ],
       [ 0.88120073],
       [-0.2978969 ]], dtype=float32), 'z_b_cohort_zerosum__': array([[ 0.08365148],
       [ 0.4037446 ],
       [ 0.219702  ],
       [-0.04724727],
       [-0.39688712],
       [-0.05613985],
       [ 0.8531689 ],
       [ 0.21144761],
       [-0.7068673 ],
       [ 0.0072215 ],
       [-0.5336951 ],
       [ 0.34807727],
       [ 0.39438143],
       [ 0.7545098 ],
       [ 0.8487969 ],
       [-0.93471295],
       [ 0.36905304],
       [ 0.23143601],
       [-0.85752094],
       [-0.93718106]], dtype=float32), 'mu_b_pos_con': array([-2.4778757, -0.7670131, -2.4533374, -2.1609292, -1.4259497],
      dtype=float32), 'z_b_pos_con_state_zerosum__': array([[-0.12270445,  0.24486493,  0.39655864, -0.2905437 , -0.34443325,
         0.00207796],
       [-0.6199098 , -0.9983607 ,  0.17514546,  0.70688826, -0.32889152,
        -0.14889538],
       [-0.7383829 , -0.3004701 , -0.3937158 ,  0.05630853,  0.04613764,
        -0.7968154 ],
       [-0.83736265,  0.53274584,  0.00774734,  0.3415022 , -0.02279032,
        -0.80878764],
       [-0.29654893,  0.8825508 , -0.31042358, -0.4510256 , -0.8262907 ,
        -0.539443  ]], dtype=float32), 'z_b_pos_con_age_zerosum__': array([[ 0.78295594,  0.05425891, -0.34404022,  0.54096305,  0.30847767],
       [-0.73004913, -0.9038148 , -0.16850933, -0.00236369, -0.6851207 ],
       [ 0.39292192,  0.9246333 ,  0.6826134 , -0.11111186,  0.9510223 ],
       [-0.83537763,  0.82692575, -0.73304725,  0.18374918, -0.9958687 ],
       [ 0.56813633,  0.7769569 , -0.94843704,  0.06147736, -0.7710584 ]],
      dtype=float32), 'z_b_pos_con_brand_zerosum__': array([[-0.9075556 ],
       [-0.27464116],
       [ 0.7783447 ],
       [-0.8754568 ],
       [-0.9165574 ]], dtype=float32), 'z_b_pos_con_cohort_zerosum__': array([[-0.37049973],
       [ 0.95798755],
       [ 0.71250516],
       [ 0.83639866],
       [-0.19618082]], dtype=float32), 'mu_b_neg_con': array([-1.0723019], dtype=float32), 'z_b_neg_con_state_zerosum__': array([[ 0.9188992 , -0.9298911 ,  0.19285995, -0.25151858,  0.37471578,
        -0.33717418]], dtype=float32), 'z_b_neg_con_age_zerosum__': array([[-0.80139756,  0.6711084 ,  0.80932933, -0.03985522,  0.22147077]],
      dtype=float32), 'z_b_neg_con_brand_zerosum__': array([[0.5175842]], dtype=float32), 'z_b_neg_con_cohort_zerosum__': array([[-0.61060673]], dtype=float32), 'mu_b_lag': array([-2.16711], dtype=float32), 'z_b_lag_state_zerosum__': array([[-0.79260015, -0.74764526, -0.82469386, -0.03429089, -0.9838945 ,
        -0.7090971 ]], dtype=float32), 'z_b_lag_age_zerosum__': array([[-0.48504904,  0.87696105, -0.742052  ,  0.9109938 ,  0.04824264]],
      dtype=float32), 'z_b_lag_brand_zerosum__': array([[0.00786572]], dtype=float32), 'z_b_lag_cohort_zerosum__': array([[-0.23193653]], dtype=float32), 'mu_b_fourier_year': array([-0.77838343,  0.6178618 ,  0.57258415,  0.7673871 ,  0.3929526 ,
        0.8627489 ,  0.4622131 , -0.769917  ,  0.13717452, -0.13153929,
       -0.83192456,  0.1140356 ,  0.05431605, -0.10188404,  0.9689602 ,
        0.05997486,  0.06146386, -0.9775481 , -0.47403213, -0.6081761 ],
      dtype=float32), 'sd_y_log__': array(3.3264434, dtype=float32)}

The variable z_a_state has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [7] [id B] <Vector(int64, shape=(1,))>
2: [] [id C] <Vector(int64, shape=(0,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [7]
2: []
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axes=None}.0, Composite{...}.1)
Toposort index: 14
Inputs types: [TensorType(float32, shape=()), TensorType(bool, shape=())]
Inputs shapes: [(), ()]
Inputs strides: [(), ()]
Inputs values: [array(-57.60402, dtype=float32), array(False)]
Outputs clients: [[output[0](z_a_state_zerosum___logprob)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],
  File "/root/.vscode-server/extensions/ms-python.debugpy-2024.12.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py", line 267, in eval_in_context
    result = eval(compiled, global_vars, local_vars)
  File "<string>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/pymc/model/core.py", line 1886, in debug
    self.value_vars, self.logp(vars=rv, sum=False)[0], on_unused_input="ignore"
  File "/usr/local/lib/python3.10/dist-packages/pymc/model/core.py", line 696, in logp
    rv_logps = transformed_conditional_logp(
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 595, in transformed_conditional_logp
    temp_logp_terms = conditional_logp(
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 122, in transformed_value_logprob
    logprobs_jac.append(logp + log_jac_det)

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_a_age has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [6] [id B] <Vector(int64, shape=(1,))>
2: [] [id C] <Vector(int64, shape=(0,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [6]
2: []
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axes=None}.0, Composite{...}.1)
Toposort index: 14
Inputs types: [TensorType(float32, shape=()), TensorType(bool, shape=())]
Inputs shapes: [(), ()]
Inputs strides: [(), ()]
Inputs values: [array(-78.23544, dtype=float32), array(False)]
Outputs clients: [[output[0](z_a_age_zerosum___logprob)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],
  File "/root/.vscode-server/extensions/ms-python.debugpy-2024.12.0-linux-x64/bundled/libs/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py", line 267, in eval_in_context
    result = eval(compiled, global_vars, local_vars)
  File "<string>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/pymc/model/core.py", line 1886, in debug
    self.value_vars, self.logp(vars=rv, sum=False)[0], on_unused_input="ignore"
  File "/usr/local/lib/python3.10/dist-packages/pymc/model/core.py", line 696, in logp
    rv_logps = transformed_conditional_logp(
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 595, in transformed_conditional_logp
    temp_logp_terms = conditional_logp(
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 122, in transformed_value_logprob
    logprobs_jac.append(logp + log_jac_det)

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_state has the following parameters:
0: 0.05 [id A] <Scalar(float32, shape=())>
1: [7] [id B] <Vector(int64, shape=(1,))>
2: [20] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.05000000074505806
1: [7]
2: [20]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(20,), ()]
Inputs strides: [(4,), ()]
Inputs values: ['not shown', array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_age has the following parameters:
0: 0.05 [id A] <Scalar(float32, shape=())>
1: [6] [id B] <Vector(int64, shape=(1,))>
2: [20] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.05000000074505806
1: [6]
2: [20]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(20,), ()]
Inputs strides: [(4,), ()]
Inputs values: ['not shown', array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_brand has the following parameters:
0: 0.05 [id A] <Scalar(float32, shape=())>
1: [2] [id B] <Vector(int64, shape=(1,))>
2: [20] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.05000000074505806
1: [2]
2: [20]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(20,), ()]
Inputs strides: [(4,), ()]
Inputs values: ['not shown', array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_cohort has the following parameters:
0: 0.05 [id A] <Scalar(float32, shape=())>
1: [2] [id B] <Vector(int64, shape=(1,))>
2: [20] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.05000000074505806
1: [2]
2: [20]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(20,), ()]
Inputs strides: [(4,), ()]
Inputs values: ['not shown', array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_pos_con_state has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [7] [id B] <Vector(int64, shape=(1,))>
2: [5] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [7]
2: [5]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(5,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([-13.464529, -93.78406 , -63.23401 , -79.51486 , -98.71708 ],
      dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_pos_con_age has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [6] [id B] <Vector(int64, shape=(1,))>
2: [5] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [6]
2: [5]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(5,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([ -49.188126,  -85.46398 , -112.68599 , -140.30869 , -114.29598 ],
      dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_pos_con_brand has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [2] [id B] <Vector(int64, shape=(1,))>
2: [5] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [2]
2: [5]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(5,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([-39.79921  ,  -2.3877418, -28.907381 , -36.937584 , -40.620224 ],
      dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_pos_con_cohort has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [2] [id B] <Vector(int64, shape=(1,))>
2: [5] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [2]
2: [5]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(5,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([ -5.4798565 , -44.503365  , -23.999535  , -33.594482  ,
        -0.54069936], dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_neg_con_state has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [7] [id B] <Vector(int64, shape=(1,))>
2: [1] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [7]
2: [1]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(1,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([-94.879524], dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_neg_con_age has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [6] [id B] <Vector(int64, shape=(1,))>
2: [1] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [6]
2: [1]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(1,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([-82.99558], dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_neg_con_cohort has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [2] [id B] <Vector(int64, shape=(1,))>
2: [1] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [2]
2: [1]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(1,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([-17.258383], dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_lag_state has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [7] [id B] <Vector(int64, shape=(1,))>
2: [1] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [7]
2: [1]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(1,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([-158.66568], dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
The variable z_b_lag_age has the following parameters:
0: 0.1 [id A] <Scalar(float32, shape=())>
1: [6] [id B] <Vector(int64, shape=(1,))>
2: [1] [id C] <Vector(int64, shape=(1,))>
The parameters evaluate to:
0: 0.10000000149011612
1: [6]
2: [1]
This does not respect one of the following constraints: mean(value, axis=n_zerosum_axes) = 0

mean(value, axis=n_zerosum_axes) = 0
Apply node that caused the error: Check{mean(value, axis=n_zerosum_axes) = 0}(Sum{axis=1}.0, All{axes=None}.0)
Toposort index: 23
Inputs types: [TensorType(float32, shape=(None,)), TensorType(bool, shape=())]
Inputs shapes: [(1,), ()]
Inputs strides: [(4,), ()]
Inputs values: [array([-112.442345], dtype=float32), array(False)]
Outputs clients: [[Sum{axes=None}(Check{mean(value, axis=n_zerosum_axes) = 0}.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/basic.py", line 529, in conditional_logp
    node_logprobs = _logprob(
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/logprob/transform_value.py", line 84, in transformed_value_logprob
    logprobs = _logprob(rv_op, values, *rv_inputs, **kwargs)
  File "/usr/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/multivariate.py", line 2841, in zerosumnormal_logp
    return check_parameters(out, *zerosums, msg="mean(value, axis=n_zerosum_axes) = 0")
  File "/usr/local/lib/python3.10/dist-packages/pymc/distributions/dist_math.py", line 74, in check_parameters
    return CheckParameterValue(msg, can_be_replaced_by_ninf)(expr, all_true_scalar)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 293, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 97, in make_node
    [value.type()],

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
```
Your traceback doesn't look like pymc=5.19, the code doesn't use `check_start_vals` inside `init_nuts` anymore. PyMC-Marketing is also specifically not compatible with more recent versions of pymc so you have to check what versions are actually running in your script. Try to print `pymc.__version__`.
Your traceback does not correspond to the example code either, it's going through a jax sampler, but I don't see it specified in `pm.sample`. Can you provide a fully reproducible snippet?
I can reproduce it though. Perhaps the logp/dlogp underflows in float32
5.19 definitely uses `check_start_vals`: https://github.com/pymc-devs/pymc/blob/a714b2485769b358195e1a1fd5bb69a337fcf526/pymc/sampling/mcmc.py#L1390
As I mentioned in the original comment, the traceback I provided is from my own model, but the reproducible example is using the radon dataset.

The only code we use from `pymc-marketing` is https://github.com/pymc-labs/pymc-marketing/blob/c8121be81d2d2b73ed5d4bf3d32f6592e3ba2afa/pymc_marketing/mmm/transformers.py#L154 and https://github.com/pymc-labs/pymc-marketing/blob/c8121be81d2d2b73ed5d4bf3d32f6592e3ba2afa/pymc_marketing/mmm/transformers.py#L433, which we could easily vendor to remove the dependency. The reason we specify `pymc-marketing==0.6.0` is specifically to avoid the restriction on using newer versions of pymc. 
Ah my bad, it does after it fails with the jitter. Either way when you sample with the pymc nuts, it's not failing in the jitter but only inside NUTS already. The `logp_dlogp_func` returns -inf for the logp at the point of failure, whereas the model.logp evaluated at that point does not. So some optimization / numerical issue is what I suspect.

Setting an older pymc-marketing gets around the pymc limitation but it's not a solution. It's just because that older version did not pin pymc as strictly as the new one does. Either way that's not the problem you're seeing here
Right, the pymc-marketing is irrelevant here.
Ah, I can see the confusion though because the "reproducible" example doesn't reproduce exactly what I was reporting. I took my coworker's word for it at 4pm on a friday afternoon haha.

This for example doesn't throw the error:
``` python
import numpy as np
import pandas as pd
import pymc as pm
import pytensor as pt

pt.config.floatX = "float32"
pt.config.warn_float64 = "ignore"

# Load the radon dataset
data = pd.read_csv(pm.get_data("radon.csv"))
data["log_radon"] = data["log_radon"].astype(np.float64)
county_idx, counties = pd.factorize(data.county)
coords = {"county": counties, "obs_id": np.arange(len(county_idx))}

# Create a simple hierarchical model for the radon dataset
with pm.Model(coords=coords, check_bounds=False) as model:
    intercept = pm.Normal("intercept", sigma=10)

    # County effects
    raw = pm.ZeroSumNormal("county_raw", dims="county")
    sd = pm.HalfNormal("county_sd")
    county_effect = pm.Deterministic("county_effect", raw * sd, dims="county")

    # Global floor effect
    floor_effect = pm.Normal("floor_effect", sigma=2)

    # County:floor interaction
    raw = pm.ZeroSumNormal("county_floor_raw", dims="county")
    sd = pm.HalfNormal("county_floor_sd")
    county_floor_effect = pm.Deterministic("county_floor_effect", raw * sd, dims="county")

    mu = (
        intercept
        + county_effect[county_idx]
        + floor_effect * data.floor.values
        + county_floor_effect[county_idx] * data.floor.values
    )

    sigma = pm.HalfNormal("sigma", sigma=1.5)
    pm.Normal("log_radon", mu=mu, sigma=sigma, observed=data.log_radon.values, dims="obs_id")

idata = pm.sample(
    model=model,
    chains=1,
    tune=500,
    draws=500,
    progressbar=False,
    compute_convergence_checks=False,
    return_inferencedata=False,
    nuts_sampler="numpyro",
    # compile_kwargs=dict(mode="NUMBA")
)
```

Commit message: removed type_checking to fix sphinx-autodoc

Patch to review:
diff --git a/pymc/initial_point.py b/pymc/initial_point.py
index 241409f683..da55fc3810 100644
--- a/pymc/initial_point.py
+++ b/pymc/initial_point.py
@@ -26,6 +26,7 @@
 
 from pymc.logprob.transforms import Transform
 from pymc.pytensorf import (
+    SeedSequenceSeed,
     compile,
     find_rng_nodes,
     replace_rng_nodes,
@@ -67,7 +68,7 @@ def make_initial_point_fns_per_chain(
     overrides: StartDict | Sequence[StartDict | None] | None,
     jitter_rvs: set[TensorVariable] | None = None,
     chains: int,
-) -> list[Callable]:
+) -> list[Callable[[SeedSequenceSeed], PointType]]:
     """Create an initial point function for each chain, as defined by initvals.
 
     If a single initval dictionary is passed, the function is replicated for each
@@ -82,6 +83,11 @@ def make_initial_point_fns_per_chain(
         Random variable tensors for which U(-1, 1) jitter shall be applied.
         (To the transformed space if applicable.)
 
+    Returns
+    -------
+    ipfns : list[Callable[[SeedSequenceSeed], dict[str, np.ndarray]]]
+        list of functions that return initial points for each chain.
+
     Raises
     ------
     ValueError
@@ -124,7 +130,7 @@ def make_initial_point_fn(
     jitter_rvs: set[TensorVariable] | None = None,
     default_strategy: str = "support_point",
     return_transformed: bool = True,
-) -> Callable:
+) -> Callable[[SeedSequenceSeed], PointType]:
     """Create seeded function that computes initial values for all free model variables.
 
     Parameters
@@ -138,6 +144,10 @@ def make_initial_point_fn(
         Initial value (strategies) to use instead of what's specified in `Model.initial_values`.
     return_transformed : bool
         If `True` the returned variables will correspond to transformed initial values.
+
+    Returns
+    -------
+    initial_point_fn : Callable[[SeedSequenceSeed], dict[str, np.ndarray]]
     """
     sdict_overrides = convert_str_to_rv_dict(model, overrides or {})
     initval_strats = {
diff --git a/pymc/model/core.py b/pymc/model/core.py
index 99711e566e..9b6c48506d 100644
--- a/pymc/model/core.py
+++ b/pymc/model/core.py
@@ -19,7 +19,7 @@
 import types
 import warnings
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from typing import (
     Literal,
     cast,
@@ -585,7 +585,7 @@ def compile_logp(
         jacobian: bool = True,
         sum: bool = True,
         **compile_kwargs,
-    ) -> PointFunc:
+    ) -> Callable[[PointType], np.ndarray]:
         """Compiled log probability density function.
 
         Parameters
diff --git a/pymc/sampling/jax.py b/pymc/sampling/jax.py
index 43e1baa87f..5226742dbc 100644
--- a/pymc/sampling/jax.py
+++ b/pymc/sampling/jax.py
@@ -18,6 +18,7 @@
 from collections.abc import Callable, Sequence
 from datetime import datetime
 from functools import partial
+from types import ModuleType
 from typing import Any, Literal
 
 import arviz as az
@@ -28,6 +29,7 @@
 
 from arviz.data.base import make_attrs
 from jax.lax import scan
+from numpy.typing import ArrayLike
 from pytensor.compile import SharedVariable, Supervisor, mode
 from pytensor.graph.basic import graph_inputs
 from pytensor.graph.fg import FunctionGraph
@@ -121,7 +123,7 @@ def _replace_shared_variables(graph: list[TensorVariable]) -> list[TensorVariabl
 def get_jaxified_graph(
     inputs: list[TensorVariable] | None = None,
     outputs: list[TensorVariable] | None = None,
-) -> list[TensorVariable]:
+) -> Callable[[list[TensorVariable]], list[TensorVariable]]:
     """Compile a PyTensor graph into an optimized JAX function."""
     graph = _replace_shared_variables(outputs) if outputs is not None else None
 
@@ -144,13 +146,13 @@ def get_jaxified_graph(
     return jax_funcify(fgraph)
 
 
-def get_jaxified_logp(model: Model, negative_logp=True) -> Callable:
+def get_jaxified_logp(model: Model, negative_logp: bool = True) -> Callable[[ArrayLike], jax.Array]:
     model_logp = model.logp()
     if not negative_logp:
         model_logp = -model_logp
     logp_fn = get_jaxified_graph(inputs=model.value_vars, outputs=[model_logp])
 
-    def logp_fn_wrap(x):
+    def logp_fn_wrap(x: ArrayLike) -> jax.Array:
         return logp_fn(*x)[0]
 
     return logp_fn_wrap
@@ -211,23 +213,43 @@ def _get_batched_jittered_initial_points(
     chains: int,
     initvals: StartDict | Sequence[StartDict | None] | None,
     random_seed: RandomSeed,
+    logp_fn: Callable[[ArrayLike], jax.Array] | None = None,
     jitter: bool = True,
     jitter_max_retries: int = 10,
 ) -> np.ndarray | list[np.ndarray]:
-    """Get jittered initial point in format expected by NumPyro MCMC kernel.
+    """Get jittered initial point in format expected by Jax MCMC kernel.
+
+    Parameters
+    ----------
+        logp_fn : Callable[Sequence[np.ndarray]], np.ndarray]
+            Jaxified logp function
 
     Returns
     -------
-    out: list of ndarrays
+    out: list[np.ndarray]
         list with one item per variable and number of chains as batch dimension.
         Each item has shape `(chains, *var.shape)`
     """
+    if logp_fn is None:
+        eval_logp_initial_point = None
+
+    else:
+
+        def eval_logp_initial_point(point: dict[str, np.ndarray]) -> jax.Array:
+            """Wrap logp_fn to conform to _init_jitter logic.
+
+            Wraps jaxified logp function to accept a dict of
+            {model_variable: np.array} key:value pairs.
+            """
+            return logp_fn(point.values())
+
     initial_points = _init_jitter(
         model,
         initvals,
         seeds=_get_seeds_per_chain(random_seed, chains),
         jitter=jitter,
         jitter_max_retries=jitter_max_retries,
+        logp_fn=eval_logp_initial_point,
     )
     initial_points_values = [list(initial_point.values()) for initial_point in initial_points]
     if chains == 1:
@@ -236,7 +258,7 @@ def _get_batched_jittered_initial_points(
 
 
 def _blackjax_inference_loop(
-    seed, init_position, logprob_fn, draws, tune, target_accept, **adaptation_kwargs
+    seed, init_position, logp_fn, draws, tune, target_accept, **adaptation_kwargs
 ):
     import blackjax
 
@@ -252,13 +274,13 @@ def _blackjax_inference_loop(
 
     adapt = blackjax.window_adaptation(
         algorithm=algorithm,
-        logdensity_fn=logprob_fn,
+        logdensity_fn=logp_fn,
         target_acceptance_rate=target_accept,
         adaptation_info_fn=get_filter_adapt_info_fn(),
         **adaptation_kwargs,
     )
     (last_state, tuned_params), _ = adapt.run(seed, init_position, num_steps=tune)
-    kernel = algorithm(logprob_fn, **tuned_params).step
+    kernel = algorithm(logp_fn, **tuned_params).step
 
     def _one_step(state, xs):
         _, rng_key = xs
@@ -289,49 +311,48 @@ def _sample_blackjax_nuts(
     tune: int,
     draws: int,
     chains: int,
-    chain_method: str | None,
+    chain_method: Literal["parallel", "vectorized"],
     progressbar: bool,
     random_seed: int,
-    initial_points,
+    initial_points: np.ndarray | list[np.ndarray],
     nuts_kwargs,
-) -> az.InferenceData:
+    logp_fn: Callable[[ArrayLike], jax.Array] | None = None,
+) -> tuple[Any, dict[str, Any], ModuleType]:
     """
     Draw samples from the posterior using the NUTS method from the ``blackjax`` library.
 
     Parameters
     ----------
-    draws : int, default 1000
-        The number of samples to draw. The number of tuned samples are discarded by
-        default.
+    model : Model, optional
+        Model to sample from. The model needs to have free random variables. When inside
+        a ``with`` model context, it defaults to that model, otherwise the model must be
+        passed explicitly.
+    target_accept : float in [0, 1].
+        The step size is tuned such that we approximate this acceptance rate. Higher
+        values like 0.9 or 0.95 often work better for problematic posteriors.
     tune : int, default 1000
         Number of iterations to tune. Samplers adjust the step sizes, scalings or
         similar during tuning. Tuning samples will be drawn in addition to the number
         specified in the ``draws`` argument.
+    draws : int, default 1000
+        The number of samples to draw. The number of tuned samples are discarded by
+        default.
     chains : int, default 4
         The number of chains to sample.
-    target_accept : float in [0, 1].
-        The step size is tuned such that we approximate this acceptance rate. Higher
-        values like 0.9 or 0.95 often work better for problematic posteriors.
+    chain_method : str, default "parallel"
+        Specify how samples should be drawn. The choices include "parallel", and
+        "vectorized".
+    progressbar : bool
+        Whether to show progressbar or not during sampling.
     random_seed : int, RandomState or Generator, optional
         Random seed used by the sampling steps.
-    initvals: StartDict or Sequence[Optional[StartDict]], optional
-        Initial values for random variables provided as a dictionary (or sequence of
-        dictionaries) mapping the random variable (by name or reference) to desired
-        starting values.
-    jitter: bool, default True
-        If True, add jitter to initial points.
-    model : Model, optional
-        Model to sample from. The model needs to have free random variables. When inside
-        a ``with`` model context, it defaults to that model, otherwise the model must be
-        passed explicitly.
+    initial_points : np.ndarray | list[np.ndarray]
+        Initial point(s) for sampler to begin sampling from.
     var_names : sequence of str, optional
         Names of variables for which to compute the posterior samples. Defaults to all
         variables in the posterior.
     keep_untransformed : bool, default False
         Include untransformed variables in the posterior samples. Defaults to False.
-    chain_method : str, default "parallel"
-        Specify how samples should be drawn. The choices include "parallel", and
-        "vectorized".
     postprocessing_backend: Optional[Literal["cpu", "gpu"]], default None,
         Specify how postprocessing should be computed. gpu or cpu
     postprocessing_vectorize: Literal["vmap", "scan"], default "scan"
@@ -343,13 +364,17 @@ def _sample_blackjax_nuts(
         ``observed_data``, ``constant_data``, ``coords``, and ``dims`` are inferred from
         the ``model`` argument if not provided in ``idata_kwargs``. If ``coords`` and
         ``dims`` are provided, they are used to update the inferred dictionaries.
+    logp_fn : Callable[[ArrayLike], jax.Array] | None:
+        jaxified logp function. If not passed in it will compute it here.
 
     Returns
     -------
-    InferenceData
-        ArviZ ``InferenceData`` object that contains the posterior samples, together
-        with their respective sample stats and pointwise log likeihood values (unless
-        skipped with ``idata_kwargs``).
+    Tuple containing:
+        raw_mcmc_samples
+            Datastructure containing raw mcmc samples
+        sample_stats : dict[str, Any]
+            Dictionary containing sample stats
+        blackjax : ModuleType["blackjax"]
     """
     import blackjax
 
@@ -366,7 +391,8 @@ def _sample_blackjax_nuts(
     if chains == 1:
         initial_points = [np.stack(init_state) for init_state in zip(initial_points)]
 
-    logprob_fn = get_jaxified_logp(model)
+    if logp_fn is None:
+        logp_fn = get_jaxified_logp(model)
 
     seed = jax.random.PRNGKey(random_seed)
     keys = jax.random.split(seed, chains)
@@ -374,7 +400,7 @@ def _sample_blackjax_nuts(
     nuts_kwargs["progress_bar"] = progressbar
     get_posterior_samples = partial(
         _blackjax_inference_loop,
-        logprob_fn=logprob_fn,
+        logp_fn=logp_fn,
         tune=tune,
         draws=draws,
         target_accept=target_accept,
@@ -386,7 +412,7 @@ def _sample_blackjax_nuts(
 
 
 # Adopted from arviz numpyro extractor
-def _numpyro_stats_to_dict(posterior):
+def _numpyro_stats_to_dict(posterior) -> dict[str, Any]:
     """Extract sample_stats from NumPyro posterior."""
     rename_key = {
         "potential_energy": "lp",
@@ -415,14 +441,58 @@ def _sample_numpyro_nuts(
     chain_method: str | None,
     progressbar: bool,
     random_seed: int,
-    initial_points,
+    initial_points: np.ndarray | list[np.ndarray],
     nuts_kwargs: dict[str, Any],
-):
+    logp_fn: Callable[[ArrayLike], jax.Array] | None = None,
+) -> tuple[Any, dict[str, Any], ModuleType]:
+    """
+    Draw samples from the posterior using the NUTS method from the ``numpyro`` library.
+
+    Parameters
+    ----------
+    model : Model, optional
+        Model to sample from. The model needs to have free random variables. When inside
+        a ``with`` model context, it defaults to that model, otherwise the model must be
+        passed explicitly.
+    target_accept : float in [0, 1].
+        The step size is tuned such that we approximate this acceptance rate. Higher
+        values like 0.9 or 0.95 often work better for problematic posteriors.
+    tune : int, default 1000
+        Number of iterations to tune. Samplers adjust the step sizes, scalings or
+        similar during tuning. Tuning samples will be drawn in addition to the number
+        specified in the ``draws`` argument.
+    draws : int, default 1000
+        The number of samples to draw. The number of tuned samples are discarded by
+        default.
+    chains : int, default 4
+        The number of chains to sample.
+    chain_method : str, default "parallel"
+        Specify how samples should be drawn. The choices include "parallel", and
+        "vectorized".
+    progressbar : bool
+        Whether to show progressbar or not during sampling.
+    random_seed : int, RandomState or Generator, optional
+        Random seed used by the sampling steps.
+    initial_points : np.ndarray | list[np.ndarray]
+        Initial point(s) for sampler to begin sampling from.
+    logp_fn : Callable[[ArrayLike], jax.Array] | None:
+        jaxified logp function. If not passed in it will compute it here.
+
+    Returns
+    -------
+    Tuple containing:
+        raw_mcmc_samples
+            Datastructure containing raw mcmc samples
+        sample_stats : dict[str, Any]
+            Dictionary containing sample stats
+        numpyro : ModuleType["numpyro"]
+    """
     import numpyro
 
     from numpyro.infer import MCMC, NUTS
 
-    logp_fn = get_jaxified_logp(model, negative_logp=False)
+    if logp_fn is None:
+        logp_fn = get_jaxified_logp(model, negative_logp=False)
 
     nuts_kwargs.setdefault("adapt_step_size", True)
     nuts_kwargs.setdefault("adapt_mass_matrix", True)
@@ -480,7 +550,7 @@ def sample_jax_nuts(
     nuts_kwargs: dict | None = None,
     progressbar: bool = True,
     keep_untransformed: bool = False,
-    chain_method: str = "parallel",
+    chain_method: Literal["parallel", "vectorized"] = "parallel",
     postprocessing_backend: Literal["cpu", "gpu"] | None = None,
     postprocessing_vectorize: Literal["vmap", "scan"] | None = None,
     postprocessing_chunks=None,
@@ -526,7 +596,7 @@ def sample_jax_nuts(
         If True, display a progressbar while sampling
     keep_untransformed : bool, default False
         Include untransformed variables in the posterior samples.
-    chain_method : str, default "parallel"
+    chain_method : Literal["parallel", "vectorized"], default "parallel"
         Specify how samples should be drawn. The choices include "parallel", and
         "vectorized".
     postprocessing_backend : Optional[Literal["cpu", "gpu"]], default None,
@@ -590,6 +660,15 @@ def sample_jax_nuts(
         get_default_varnames(filtered_var_names, include_transformed=keep_untransformed)
     )
 
+    if nuts_sampler == "numpyro":
+        sampler_fn = _sample_numpyro_nuts
+        logp_fn = get_jaxified_logp(model, negative_logp=False)
+    elif nuts_sampler == "blackjax":
+        sampler_fn = _sample_blackjax_nuts
+        logp_fn = get_jaxified_logp(model)
+    else:
+        raise ValueError(f"{nuts_sampler=} not recognized")
+
     (random_seed,) = _get_seeds_per_chain(random_seed, 1)
 
     initial_points = _get_batched_jittered_initial_points(
@@ -598,15 +677,9 @@ def sample_jax_nuts(
         initvals=initvals,
         random_seed=random_seed,
         jitter=jitter,
+        logp_fn=logp_fn,
     )
 
-    if nuts_sampler == "numpyro":
-        sampler_fn = _sample_numpyro_nuts
-    elif nuts_sampler == "blackjax":
-        sampler_fn = _sample_blackjax_nuts
-    else:
-        raise ValueError(f"{nuts_sampler=} not recognized")
-
     tic1 = datetime.now()
     raw_mcmc_samples, sample_stats, library = sampler_fn(
         model=model,
diff --git a/pymc/sampling/mcmc.py b/pymc/sampling/mcmc.py
index bc3e3475d1..c229469992 100644
--- a/pymc/sampling/mcmc.py
+++ b/pymc/sampling/mcmc.py
@@ -1338,7 +1338,7 @@ def _init_jitter(
     seeds: Sequence[int] | np.ndarray,
     jitter: bool,
     jitter_max_retries: int,
-    logp_dlogp_func=None,
+    logp_fn: Callable[[PointType], np.ndarray] | None = None,
 ) -> list[PointType]:
     """Apply a uniform jitter in [-1, 1] to the test value as starting point in each chain.
 
@@ -1353,11 +1353,14 @@ def _init_jitter(
         Whether to apply jitter or not.
     jitter_max_retries : int
         Maximum number of repeated attempts at initializing values (per chain).
+    logp_fn: Callable[[dict[str, np.ndarray]], np.ndarray | jax.Array] | None
+        logp function that takes the output of initial point functions as input.
+        If None, will use the results of model.compile_logp().
 
     Returns
     -------
-    start : ``pymc.model.Point``
-        Starting point for sampler
+    initial_points : list[dict[str, np.ndarray]]
+        List of starting points for the sampler
     """
     ipfns = make_initial_point_fns_per_chain(
         model=model,
@@ -1369,14 +1372,10 @@ def _init_jitter(
     if not jitter:
         return [ipfn(seed) for ipfn, seed in zip(ipfns, seeds)]
 
-    model_logp_fn: Callable
-    if logp_dlogp_func is None:
+    if logp_fn is None:
         model_logp_fn = model.compile_logp()
     else:
-
-        def model_logp_fn(ip):
-            q, _ = DictToArrayBijection.map(ip)
-            return logp_dlogp_func([q], extra_vars={})[0]
+        model_logp_fn = logp_fn
 
     initial_points = []
     for ipfn, seed in zip(ipfns, seeds):
@@ -1501,13 +1500,18 @@ def init_nuts(
 
     logp_dlogp_func = model.logp_dlogp_function(ravel_inputs=True, **compile_kwargs)
     logp_dlogp_func.trust_input = True
+
+    def model_logp_fn(ip: PointType) -> np.ndarray:
+        q, _ = DictToArrayBijection.map(ip)
+        return logp_dlogp_func([q], extra_vars={})[0]
+
     initial_points = _init_jitter(
         model,
         initvals,
         seeds=random_seed_list,
         jitter="jitter" in init,
         jitter_max_retries=jitter_max_retries,
-        logp_dlogp_func=logp_dlogp_func,
+        logp_fn=model_logp_fn,
     )
 
     apoints = [DictToArrayBijection.map(point) for point in initial_points]


Code review comments:
diff_hunk: @@ -585,7 +585,7 @@ def compile_logp(
         jacobian: bool = True,
         sum: bool = True,
         **compile_kwargs,
-    ) -> PointFunc:
+    ) -> Callable[[PointType], np.ndarray]:
path: pymc/model/core.py
text: @user1:
Doesn't make sense to override compile_logp, but not others. Besides it's not true, when sum=False, it will return a list of np.ndarray. And the output is indeed a `PointFunc` class. If anything you could change the signature of `PointFunc.__call__`. And technically the output can be other types also if the compilation backend is JAX for instance (or it could be Sparse scipy matrices)

@author:
I think I changed this just to keep things straight in my own head as I was working on it, if `mcmc._init_jitter` calls `model.compile_logp()` it will always return a `Callable[[PointType], np.ndarray]`.

`PointFunc` is indeed the type that is returned but it doesn't really describe what the object does without investigating the definition of `PointFunc`. If it used a decorator pattern rather than a class it could be made a little clearer but that's probably a separate piece of work so I'll revert it here.

diff_hunk: @@ -211,23 +213,43 @@ def _get_batched_jittered_initial_points(
     chains: int,
     initvals: StartDict | Sequence[StartDict | None] | None,
     random_seed: RandomSeed,
+    logp_fn: Callable[[ArrayLike], jax.Array] | None = None,
     jitter: bool = True,
     jitter_max_retries: int = 10,
 ) -> np.ndarray | list[np.ndarray]:
-    """Get jittered initial point in format expected by NumPyro MCMC kernel.
+    """Get jittered initial point in format expected by Jax MCMC kernel.
+
+    Parameters
+    ----------
+        logp_fn : Callable[Sequence[np.ndarray]], np.ndarray]
+            Jaxified logp function
 
     Returns
     -------
-    out: list of ndarrays
+    out: list[np.ndarray]
path: pymc/sampling/jax.py
text: @user1:
IIRC we use human readable definition on the docstrings (so what was there before)

@author:
implemented

diff_hunk: @@ -289,49 +311,48 @@ def _sample_blackjax_nuts(
     tune: int,
     draws: int,
     chains: int,
-    chain_method: str | None,
+    chain_method: Literal["parallel", "vectorized"],
     progressbar: bool,
     random_seed: int,
-    initial_points,
+    initial_points: np.ndarray | list[np.ndarray],
     nuts_kwargs,
-) -> az.InferenceData:
+    logp_fn: Callable[[ArrayLike], jax.Array] | None = None,
+) -> tuple[Any, dict[str, Any], ModuleType]:
     """
     Draw samples from the posterior using the NUTS method from the ``blackjax`` library.
 
     Parameters
     ----------
-    draws : int, default 1000
-        The number of samples to draw. The number of tuned samples are discarded by
-        default.
+    model : Model, optional
+        Model to sample from. The model needs to have free random variables. When inside
+        a ``with`` model context, it defaults to that model, otherwise the model must be
+        passed explicitly.
+    target_accept : float in [0, 1].
+        The step size is tuned such that we approximate this acceptance rate. Higher
+        values like 0.9 or 0.95 often work better for problematic posteriors.
     tune : int, default 1000
         Number of iterations to tune. Samplers adjust the step sizes, scalings or
         similar during tuning. Tuning samples will be drawn in addition to the number
         specified in the ``draws`` argument.
+    draws : int, default 1000
+        The number of samples to draw. The number of tuned samples are discarded by
+        default.
     chains : int, default 4
         The number of chains to sample.
-    target_accept : float in [0, 1].
-        The step size is tuned such that we approximate this acceptance rate. Higher
-        values like 0.9 or 0.95 often work better for problematic posteriors.
+    chain_method : str, default "parallel"
path: pymc/sampling/jax.py
text: @user1:
There's no default in these inner functions?

@author:
The majority of this docstring is unchanged from the original, just moved around to match the order in the function signature. I noticed that the function doesn't actually have defaults but my thinking was that these private functions are always called as `sampler_fn` in `sample_jax_nuts` with the values being passed in from there. So the defaults of `sample_jax_nuts` are effectively the defaults of `_sample_numpyro_nuts`/`_sample_blackjax_nuts`. To save people from having to work out that connection themselves we could include it in the docstring here? No one will be manually calling this but it could help anyone who is looking into the source code. 

@author:
I added an explanatory line to the docstring that makes this clearer

@user1:
An argument against that is that it may become outdated if/when the default changes in the function that calls this. Since this is not an actual default, the developer would be forgiven for not checking the docstrings of an inner called function needed updating.

Also in what pertains to this function it's objectively wrong. You need to pass it, there's no default. This kind of work is why I am generally okay with not documenting non-user-facing functions like this, but I'm not against it if someone really wants it 

@author:
Ok, no problem. I've updated the docstrings to more accurately reflect the function signature.

diff_hunk: @@ -343,13 +364,17 @@ def _sample_blackjax_nuts(
         ``observed_data``, ``constant_data``, ``coords``, and ``dims`` are inferred from
         the ``model`` argument if not provided in ``idata_kwargs``. If ``coords`` and
         ``dims`` are provided, they are used to update the inferred dictionaries.
+    logp_fn : Callable[[ArrayLike], jax.Array] | None:
+        jaxified logp function. If not passed in it will compute it here.
path: pymc/sampling/jax.py
text: @user1:
```suggestion
    logp_fn : Callable[[ArrayLike], jax.Array] | None:
        jaxified logp function. If not passed in it will created anew.
```

@author:
implemented

diff_hunk: @@ -415,14 +441,58 @@ def _sample_numpyro_nuts(
     chain_method: str | None,
     progressbar: bool,
     random_seed: int,
-    initial_points,
+    initial_points: np.ndarray | list[np.ndarray],
     nuts_kwargs: dict[str, Any],
-):
+    logp_fn: Callable[[ArrayLike], jax.Array] | None = None,
+) -> tuple[Any, dict[str, Any], ModuleType]:
+    """
+    Draw samples from the posterior using the NUTS method from the ``numpyro`` library.
+
+    Parameters
+    ----------
+    model : Model, optional
+        Model to sample from. The model needs to have free random variables. When inside
+        a ``with`` model context, it defaults to that model, otherwise the model must be
+        passed explicitly.
+    target_accept : float in [0, 1].
+        The step size is tuned such that we approximate this acceptance rate. Higher
+        values like 0.9 or 0.95 often work better for problematic posteriors.
+    tune : int, default 1000
+        Number of iterations to tune. Samplers adjust the step sizes, scalings or
+        similar during tuning. Tuning samples will be drawn in addition to the number
+        specified in the ``draws`` argument.
+    draws : int, default 1000
+        The number of samples to draw. The number of tuned samples are discarded by
+        default.
+    chains : int, default 4
+        The number of chains to sample.
+    chain_method : str, default "parallel"
+        Specify how samples should be drawn. The choices include "parallel", and
+        "vectorized".
+    progressbar : bool
+        Whether to show progressbar or not during sampling.
+    random_seed : int, RandomState or Generator, optional
+        Random seed used by the sampling steps.
+    initial_points : np.ndarray | list[np.ndarray]
+        Initial point(s) for sampler to begin sampling from.
+    logp_fn : Callable[[ArrayLike], jax.Array] | None:
+        jaxified logp function. If not passed in it will compute it here.
+
+    Returns
+    -------
+    Tuple containing:
path: pymc/sampling/jax.py
text: @user1:
Multiple output of a python function are always a tuple, so no need to specify `Tuple containing`, just describe the 3 outputs. Same above

@author:
Implemented