# Integrating custom JAX-based likelihoods (e.g. auto differentiable stochastic filters) w/ PyMC sampling workflow

**URL:** <https://discourse.pymc.io/t/integrating-custom-jax-based-likelihoods-e-g-auto-differentiable-stochastic-filters-w-pymc-sampling-workflow/17341>\
**Category:** General\
**Tags:** jax, ode\_pde\
**Created:** [September 21, 2025, 5:54pm UTC](https://discourse.pymc.io/t/integrating-custom-jax-based-likelihoods-e-g-auto-differentiable-stochastic-filters-w-pymc-sampling-workflow/17341 "2025-09-21T17:54:45Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![mattlevine22](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/mattlevine22/32/9432_2.png) [@mattlevine22](https://discourse.pymc.io/u/mattlevine22)\
**Post date:** [September 21, 2025, 5:54pm UTC](https://discourse.pymc.io/t/integrating-custom-jax-based-likelihoods-e-g-auto-differentiable-stochastic-filters-w-pymc-sampling-workflow/17341/1 "2025-09-21T17:54:45Z")

</div>

I’m interested in Bayesian parameter inference for partially-observed stochastic differential equations, and have been building a JAX-based repo that computes the marginal log probability of given time series data under a specified parametric model (exactly for linear SDEs, approximately for non-linear SDEs…using Kalman filters and their non-linear variants). I’ve written demos for how to integrate with numpyro via a simple `numpyro.factor()`call. I’m curious what it would take to write a similar type of demo for PyMC. Can PyMC use/benefit-from JAX-based likelihood?

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [September 22, 2025, 7:53am UTC](https://discourse.pymc.io/t/integrating-custom-jax-based-likelihoods-e-g-auto-differentiable-stochastic-filters-w-pymc-sampling-workflow/17341/2 "2025-09-22T07:53:31Z")

</div>

Check [How to use JAX ODEs and Neural Networks in PyMC](https://www.pymc-labs.com/blog-posts/jax-functions-in-pymc-3-quick-examples)

And soon more easily: [Implement as\_jax\_op by aseyboldt · Pull Request #1614 · pymc-devs/pytensor · GitHub](https://github.com/pymc-devs/pytensor/pull/1614)

---

<div class="post-metadata">

**Author:** ![mattlevine22](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/mattlevine22/32/9432_2.png) [@mattlevine22](https://discourse.pymc.io/u/mattlevine22)\
**Post date:** [September 22, 2025, 2:36pm UTC](https://discourse.pymc.io/t/integrating-custom-jax-based-likelihoods-e-g-auto-differentiable-stochastic-filters-w-pymc-sampling-workflow/17341/3 "2025-09-22T14:36:45Z")

</div>

Thanks! The PR you linked looks a lot easier to work with—maybe I should wait until that is out to make my demo? Also, does this look like the correct way to use a) your as\_jax\_op syntax and b) the pm.Potential feature (sorry, I’m new to PyMC) ?

```python
import numpy as np
import jax.numpy as jnp
import pymc as pm
from pytensor.link.jax import as_jax_op

# ---- Example data ----
rng = np.random.default_rng(0)
n, p = 50, 3
X = rng.normal(size=(n, p))
true_beta = np.array([2.0, -1.0, 0.5])
y = X @ true_beta + rng.normal(scale=1.0, size=n)

# ---- Define JAX log-likelihood for regression ----
def loglike_fn(beta, X_obs, y_obs):
    mu = X_obs @ beta
    return jnp.sum(-0.5 * (y_obs - mu) ** 2)

# ---- Wrap JAX function as pytensor Op ----
@as_jax_op
def loglike(beta, X_obs, y_obs):
    return loglike_fn(beta, X_obs, y_obs)

# ---- Build and sample PyMC model ----
with pm.Model() as model:
    beta = pm.Uniform("beta", lower=-5, upper=5, shape=p) # <-- only real change!
    pm.Potential("likelihood", loglike(beta, X, y))
    idata = pm.sample_numpyro_nuts(draws=1000, tune=500)

print(idata.posterior)

```