# \#jax

**URL:** https://discourse.pymc.io/tag/jax/84.md

[Latest](https://discourse.pymc.io/latest.md) · [Categories](https://discourse.pymc.io/categories.md) · [Tags](https://discourse.pymc.io/tags.md)

---

## [Possible bug: sample\_prior\_predictive with JAX fails for time-varying statespace models](https://discourse.pymc.io/t/possible-bug-sample-prior-predictive-with-jax-fails-for-time-varying-statespace-models/17549)

<div class="topic-metadata">

**Author:** [@JSR](https://discourse.pymc.io/u/JSR)\
**Replies:** 2\
**Last updated:** [February 17, 2026, 8:37pm UTC](https://discourse.pymc.io/t/possible-bug-sample-prior-predictive-with-jax-fails-for-time-varying-statespace-models/17549 "2026-02-17T20:37:38Z")

</div>

In the code below, sample\_prior\_predictive(compile\_kwargs={“mode”: “JAX”}) fails for the time-varying model. (Simple one-state model with exponential decay and a source term which may or may not change with time.) Also,…

---

## [Integrating custom JAX-based likelihoods (e.g. auto differentiable stochastic filters) w/ PyMC sampling workflow](https://discourse.pymc.io/t/integrating-custom-jax-based-likelihoods-e-g-auto-differentiable-stochastic-filters-w-pymc-sampling-workflow/17341)

<div class="topic-metadata">

**Author:** [@mattlevine22](https://discourse.pymc.io/u/mattlevine22)\
**Replies:** 2\
**Last updated:** [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 "2025-09-22T14:36: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 …

---

## [Pymc State Space .1.4 -- XlaRuntimeError Incorrect output dtype for return value #0: Expected: int64, Actual: int32](https://discourse.pymc.io/t/pymc-state-space-1-4-xlaruntimeerror-incorrect-output-dtype-for-return-value-0-expected-int64-actual-int32/16295)

<div class="topic-metadata">

**Author:** [@rOy\_bOy](https://discourse.pymc.io/u/rOy_bOy)\
**Replies:** 10\
**Last updated:** [February 27, 2025, 9:22am UTC](https://discourse.pymc.io/t/pymc-state-space-1-4-xlaruntimeerror-incorrect-output-dtype-for-return-value-0-expected-int64-actual-int32/16295 "2025-02-27T09:22:26Z")

</div>

Hey State Space Team: I recently updated to .1.4 to check out the regression component. My original model, running on an earlier version would sample fine, but now not so much. I get the following error when I try to …

---

## [AttributeError: 'numpy.ndarray' object has no attribute 'at' when sampling LKJ Cholesky Covariance Priors for Multivariate Normal Models example with numpyro or blackjax](https://discourse.pymc.io/t/attributeerror-numpy-ndarray-object-has-no-attribute-at-when-sampling-lkj-cholesky-covariance-priors-for-multivariate-normal-models-example-with-numpyro-or-blackjax/16598)

<div class="topic-metadata">

**Author:** [@EhDings](https://discourse.pymc.io/u/EhDings)\
**Replies:** 4\
**Last updated:** [February 26, 2025, 1:53pm UTC](https://discourse.pymc.io/t/attributeerror-numpy-ndarray-object-has-no-attribute-at-when-sampling-lkj-cholesky-covariance-priors-for-multivariate-normal-models-example-with-numpyro-or-blackjax/16598 "2025-02-26T13:53:43Z")

</div>

I wanted to try out sampling with numpyro on an AMD GPU using a model based on this example: https://www.pymc.io/projects/examples/en/latest/howto/LKJ.html It works with PyMC and nutpie samplers, but not with jax sampl…

---

## [Numpyro sampler issue](https://discourse.pymc.io/t/numpyro-sampler-issue/16497)

<div class="topic-metadata">

**Author:** [@leadhopper](https://discourse.pymc.io/u/leadhopper)\
**Replies:** 1\
**Last updated:** [February 5, 2025, 6:28am UTC](https://discourse.pymc.io/t/numpyro-sampler-issue/16497 "2025-02-05T06:28:31Z")

</div>

In executing MCMC, if you specify nuts\_sampler=“numpyro”, the following error occurs: “NotImplementedError: jax.experimental.host\_callback has been deprecated since March 2024 and is now no longer supported. See Depreca…

---

## [How to write a PyTensor Op to wrap Jax ODEs with multiple input parameters](https://discourse.pymc.io/t/how-to-write-a-pytensor-op-to-wrap-jax-odes-with-multiple-input-parameters/11611)

<div class="topic-metadata">

**Author:** [@nlinden](https://discourse.pymc.io/u/nlinden)\
**Replies:** 7\
**Last updated:** [July 28, 2024, 6:34pm UTC](https://discourse.pymc.io/t/how-to-write-a-pytensor-op-to-wrap-jax-odes-with-multiple-input-parameters/11611 "2024-07-28T18:34:55Z")

</div>

I am new to PyMC and Jax. I am trying to use PyMC to estimate the parameters for ODE implemented in Jax. The ODEs I work with often have many input parameters so I would like to write a PyTensor Op that can handle this. …

---

## [Jax-Metal pm.Sample with numpyro on M1/M2](https://discourse.pymc.io/t/jax-metal-pm-sample-with-numpyro-on-m1-m2/14591)

<div class="topic-metadata">

**Author:** [@ajc](https://discourse.pymc.io/u/ajc)\
**Replies:** 1\
**Last updated:** [June 12, 2024, 5:42am UTC](https://discourse.pymc.io/t/jax-metal-pm-sample-with-numpyro-on-m1-m2/14591 "2024-06-12T05:42:15Z")

</div>

Curious if anyone has been able to successfully sample using Jax-Metal on M1/M2 Apple Silicon? So far I’ve only been able to trade an XLA runtime error for a segmentation fault.

---

## [New distribution for PyMC: alpha-stable family \[beta\]](https://discourse.pymc.io/t/new-distribution-for-pymc-alpha-stable-family-beta/14398)

<div class="topic-metadata">

**Author:** [@tjhunter](https://discourse.pymc.io/u/tjhunter)\
**Replies:** 5\
**Last updated:** [May 28, 2024, 8:51am UTC](https://discourse.pymc.io/t/new-distribution-for-pymc-alpha-stable-family-beta/14398 "2024-05-28T08:51:50Z")

</div>

For the folks who need robust estimation, I am excited to release a first public version of levy-stable-jax, a package for 1-dimensional alpha-stable distributions. These distributions generalize Gaussian distributions t…

---

## [Recent error from Jax's (numpy) clip method](https://discourse.pymc.io/t/recent-error-from-jaxs-numpy-clip-method/14457)

<div class="topic-metadata">

**Author:** [@Allan\_Campopiano](https://discourse.pymc.io/u/Allan_Campopiano)\
**Replies:** 2\
**Last updated:** [May 21, 2024, 3:51pm UTC](https://discourse.pymc.io/t/recent-error-from-jaxs-numpy-clip-method/14457 "2024-05-21T15:51:25Z")

</div>

Hi folks, Recently I began getting an error when using the blackjax NUTS sampler. For example, when executing pm.sample as follows: idata=pm.sample(draws=2000, nuts\_sampler='blackjax', step=pm.NUTS(target\_accept=0.95),…

---

## [Jax sampling for bayesian neural network](https://discourse.pymc.io/t/jax-sampling-for-bayesian-neural-network/14095)

<div class="topic-metadata">

**Author:** [@Baba\_Yara\_Fahiz](https://discourse.pymc.io/u/Baba_Yara_Fahiz)\
**Replies:** 3\
**Last updated:** [March 27, 2024, 9:37pm UTC](https://discourse.pymc.io/t/jax-sampling-for-bayesian-neural-network/14095 "2024-03-27T21:37:43Z")

</div>

I am working with a version of the bayesian neural network model in this example notebook. I was wondering how one can change this to sample from the GPU using a JAX backend?

---

## [Can jax be used to wrap a function that contained python libraries?](https://discourse.pymc.io/t/can-jax-be-used-to-wrap-a-function-that-contained-python-libraries/13489)

<div class="topic-metadata">

**Author:** [@liu\_liu](https://discourse.pymc.io/u/liu_liu)\
**Replies:** 2\
**Last updated:** [December 25, 2023, 9:25am UTC](https://discourse.pymc.io/t/can-jax-be-used-to-wrap-a-function-that-contained-python-libraries/13489 "2023-12-25T09:25:57Z")

</div>

Hi all, I am trying to wrap my custom likelihood function using jax. I followed How to use JAX ODEs and Neural Networks in PyMC guidance. My custom likelihood function depends on the cantera module. Part of my code ar…

---

## [How to resolve Input Dimension Mis-match Error in Hierarchical Bayesian Inference with PyMC3](https://discourse.pymc.io/t/how-to-resolve-input-dimension-mis-match-error-in-hierarchical-bayesian-inference-with-pymc3/13471)

<div class="topic-metadata">

**Author:** [@Aravind27](https://discourse.pymc.io/u/Aravind27)\
**Replies:** 3\
**Last updated:** [December 20, 2023, 1:47pm UTC](https://discourse.pymc.io/t/how-to-resolve-input-dimension-mis-match-error-in-hierarchical-bayesian-inference-with-pymc3/13471 "2023-12-20T13:47:30Z")

</div>

Hello everyone, I am trying to implement a slightly peculiar model using hierarchical bayesian inference with pymc3 and have been encountering a ‘ValueError’ due to input dimension mis-match in Pymc3. The model has two …

---

## [pm.Categorical with sample\_numpyro\_nuts](https://discourse.pymc.io/t/pm-categorical-with-sample-numpyro-nuts/13361)

<div class="topic-metadata">

**Author:** [@Gumibear\_Toronto](https://discourse.pymc.io/u/Gumibear_Toronto)\
**Replies:** 3\
**Last updated:** [November 24, 2023, 4:01pm UTC](https://discourse.pymc.io/t/pm-categorical-with-sample-numpyro-nuts/13361 "2023-11-24T16:01:20Z")

</div>

Hi, I have used sample\_numpyro\_nuts because the sampler is less time consuming. It even works well with pm.Categorial() with observed; however, if I want to make predictions with the new data, it poses a problem. with …

---

## [Dynamic shaping, "round" function, JAX, and a "few" more questions](https://discourse.pymc.io/t/dynamic-shaping-round-function-jax-and-a-few-more-questions/12264)

<div class="topic-metadata">

**Author:** [@caclav](https://discourse.pymc.io/u/caclav)\
**Replies:** 28\
**Last updated:** [September 29, 2023, 6:31pm UTC](https://discourse.pymc.io/t/dynamic-shaping-round-function-jax-and-a-few-more-questions/12264 "2023-09-29T18:31:13Z")

</div>

Thank you for all the hard work you have put into providing us with a powerful and flexible Bayesian inference framework in Python! Looking forward to seeing what future developments PyMC has in store for us! I am tryin…

---

## [JAX Sampling Error with TruncatedNormal Distribution](https://discourse.pymc.io/t/jax-sampling-error-with-truncatednormal-distribution/10656)

<div class="topic-metadata">

**Author:** [@jordan.howell2](https://discourse.pymc.io/u/jordan.howell2)\
**Replies:** 6\
**Last updated:** [July 27, 2023, 12:05pm UTC](https://discourse.pymc.io/t/jax-sampling-error-with-truncatednormal-distribution/10656 "2023-07-27T12:05:17Z")

</div>

Hello, Has anyone had issues with using a TruncatedNormal on the observed data and jax sampling? I changed a model from using the Normal distribution to a trunated normal and got he following error: AttributeError: mo…

---

## [Numpyro JAX sampling very slow](https://discourse.pymc.io/t/numpyro-jax-sampling-very-slow/12469)

<div class="topic-metadata">

**Author:** [@Varun\_Gupta](https://discourse.pymc.io/u/Varun_Gupta)\
**Replies:** 1\
**Last updated:** [July 9, 2023, 10:05am UTC](https://discourse.pymc.io/t/numpyro-jax-sampling-very-slow/12469 "2023-07-09T10:05:46Z")

</div>

Hi, I have a custom log likelihood function and I have vectorized all the calculations. There are no python loops or scans in the code. Numpyro JAX sampling is extremely slow, with 1-3 its/sec. I believe with large inpu…

---

## [Jax on Apple Silicon GPU](https://discourse.pymc.io/t/jax-on-apple-silicon-gpu/12300)

<div class="topic-metadata">

**Author:** [@zweli](https://discourse.pymc.io/u/zweli)\
**Replies:** 2\
**Last updated:** [June 13, 2023, 6:36am UTC](https://discourse.pymc.io/t/jax-on-apple-silicon-gpu/12300 "2023-06-13T06:36:42Z")

</div>

Jax GPU access for us Mac users with the new M1/M2 silicon has been non-existent. This might change with the recent release by Apple of jax-metal. I have begun trying to get it to work but wanted to post this here in ca…

---

## [NameError: unbound axis name raised during transformation of variables after sample\_numpyro\_nuts](https://discourse.pymc.io/t/nameerror-unbound-axis-name-raised-during-transformation-of-variables-after-sample-numpyro-nuts/11167)

<div class="topic-metadata">

**Author:** [@n1ck94](https://discourse.pymc.io/u/n1ck94)\
**Replies:** 6\
**Last updated:** [June 13, 2023, 2:49am UTC](https://discourse.pymc.io/t/nameerror-unbound-axis-name-raised-during-transformation-of-variables-after-sample-numpyro-nuts/11167 "2023-06-13T02:49:46Z")

</div>

Hello, Has anyone come across an error like this before? If so, any advice would be appreciated. My likelihood function is a jax function which is then wrapped in PyTensor and subsequently unwrapped so it can be sample…

---

## [Wild results for masked multinomial with jax sampler(s)](https://discourse.pymc.io/t/wild-results-for-masked-multinomial-with-jax-sampler-s/12067)

<div class="topic-metadata">

**Author:** [@jason](https://discourse.pymc.io/u/jason)\
**Replies:** 2\
**Last updated:** [May 7, 2023, 5:16am UTC](https://discourse.pymc.io/t/wild-results-for-masked-multinomial-with-jax-sampler-s/12067 "2023-05-07T05:16:17Z")

</div>

Hi, I’m working on a multinomial model whereby not all choices may be available all the time. The standard pymc and native numpyro models and samples look fine. But I see strange traces when using the pymc jaxified model…

---

## [Pymc3 on GPU using jax](https://discourse.pymc.io/t/pymc3-on-gpu-using-jax/11896)

<div class="topic-metadata">

**Author:** [@Viktoria\_S](https://discourse.pymc.io/u/Viktoria_S)\
**Replies:** 2\
**Last updated:** [April 20, 2023, 8:24am UTC](https://discourse.pymc.io/t/pymc3-on-gpu-using-jax/11896 "2023-04-20T08:24:02Z")

</div>

Hi everyone, I am trying to create a conda environment using pymc3 with jax following this link. However, it gives me the following error: Collecting git+https://github.com/pymc-devs/pymc3.git@pymc3jax Cloning https://…

---

## [Out of Memory when using pm.sampling.jax.sample\_blackjax\_nuts](https://discourse.pymc.io/t/out-of-memory-when-using-pm-sampling-jax-sample-blackjax-nuts/11544)

<div class="topic-metadata">

**Author:** [@jroberayalas](https://discourse.pymc.io/u/jroberayalas)\
**Replies:** 2\
**Last updated:** [March 23, 2023, 7:49pm UTC](https://discourse.pymc.io/t/out-of-memory-when-using-pm-sampling-jax-sample-blackjax-nuts/11544 "2023-03-23T19:49:50Z")

</div>

Hi, I’ve a question regarding the usage of pm.sampling.jax.sample\_blackjax\_nuts. I’ve a training dataset with roughly 11K rows and I built a model to leverage the hierarchical structure within it. When I use pm.sample(r…

---

## [How can I output a gradient in vector format in Op.grad instance?](https://discourse.pymc.io/t/how-can-i-output-a-gradient-in-vector-format-in-op-grad-instance/10288)

<div class="topic-metadata">

**Author:** [@HajimeKawahara](https://discourse.pymc.io/u/HajimeKawahara)\
**Replies:** 7\
**Last updated:** [January 14, 2023, 1:13pm UTC](https://discourse.pymc.io/t/how-can-i-output-a-gradient-in-vector-format-in-op-grad-instance/10288 "2023-01-14T13:13:56Z")

</div>

Hello, I am new to PyMC. I’ve been using JAX+numpyro, but now I’m going to PyMC 4.0 with JAX. I am trying to wrap a JAX function into aesara while reading this page. I am having trouble specifying the gradient of a vec…

---

## [How to use JAX ODEs and Neural Networks in PyMC](https://discourse.pymc.io/t/how-to-use-jax-odes-and-neural-networks-in-pymc/11109)

<div class="topic-metadata">

**Author:** [@twiecki](https://discourse.pymc.io/u/twiecki)\
**Replies:** 0\
**Last updated:** [January 4, 2023, 3:25am UTC](https://discourse.pymc.io/t/how-to-use-jax-odes-and-neural-networks-in-pymc/11109 "2023-01-04T03:25:50Z")

</div>

"PyMC strength comes from its expressiveness. If you have a data-generating process and want to infer parameters of interest, all you need to do is write it down, choose some priors and let it sample. Sometimes this is …

---

## [Numpyro Convergence Diagnostics](https://discourse.pymc.io/t/numpyro-convergence-diagnostics/11081)

<div class="topic-metadata">

**Author:** [@Thops](https://discourse.pymc.io/u/Thops)\
**Replies:** 1\
**Last updated:** [December 30, 2022, 11:30pm UTC](https://discourse.pymc.io/t/numpyro-convergence-diagnostics/11081 "2022-12-30T23:30:14Z")

</div>

I’ve swapped to numpyro and jax from sampling, however I don’t see any mention of divergences. Does sampling.jax.sample\_numpyro\_nuts run the checks the regular pymc.sample does?

---

## [Out of memory when "transforming variables" in Numpyro & JAX](https://discourse.pymc.io/t/out-of-memory-when-transforming-variables-in-numpyro-jax/9904)

<div class="topic-metadata">

**Author:** [@jhrcook](https://discourse.pymc.io/u/jhrcook)\
**Replies:** 10\
**Last updated:** [December 12, 2022, 10:02am UTC](https://discourse.pymc.io/t/out-of-memory-when-transforming-variables-in-numpyro-jax/9904 "2022-12-12T10:02:06Z")

</div>

I am able to finish the sampling process on a large model using the JAX backend with a GPU, but it fails during the “transforming variables” step because it requires too much RAM. I set postprocessing\_backend="cpu" which…

---

## [Sample\_numpyro\_nuts hangs when parallelizing over datasets with multiprocessing](https://discourse.pymc.io/t/sample-numpyro-nuts-hangs-when-parallelizing-over-datasets-with-multiprocessing/10788)

<div class="topic-metadata">

**Author:** [@energyatinfinity](https://discourse.pymc.io/u/energyatinfinity)\
**Replies:** 0\
**Last updated:** [November 9, 2022, 6:25pm UTC](https://discourse.pymc.io/t/sample-numpyro-nuts-hangs-when-parallelizing-over-datasets-with-multiprocessing/10788 "2022-11-09T18:25:32Z")

</div>

I have a similar use case to the post linked below - fitting a pymc model to many datasets in parallel. My implementation for fitting in parallel is pretty much the same as in the linked post, but for sampling I’m tryin…

---

## [JAX Sampling Error when same model had previously worked without it](https://discourse.pymc.io/t/jax-sampling-error-when-same-model-had-previously-worked-without-it/10677)

<div class="topic-metadata">

**Author:** [@iarzt](https://discourse.pymc.io/u/iarzt)\
**Replies:** 1\
**Last updated:** [October 25, 2022, 9:39pm UTC](https://discourse.pymc.io/t/jax-sampling-error-when-same-model-had-previously-worked-without-it/10677 "2022-10-25T21:39:44Z")

</div>

I have a mixed model that uses priors with a normal distribution as well as a Bernoulli distribution. Recently, I have switched to using the JAX sampler after gaining access to a GPU. Previously, I was able to run this m…

---

## [Slow inference for numpyro sampling on Colab GPU](https://discourse.pymc.io/t/slow-inference-for-numpyro-sampling-on-colab-gpu/10453)

<div class="topic-metadata">

**Author:** [@mattiasthalen](https://discourse.pymc.io/u/mattiasthalen)\
**Replies:** 3\
**Last updated:** [October 5, 2022, 7:22am UTC](https://discourse.pymc.io/t/slow-inference-for-numpyro-sampling-on-colab-gpu/10453 "2022-10-05T07:22:01Z")

</div>

When using pm.sampling\_jax.sample\_numpyro\_nuts on Colab with GPU, my model drops to \<2 it/s, but using the same with CPU runs ~10it/s. Has anyone else noticed this?

---

## [Var\_names not working with sample\_numpyro\_nuts](https://discourse.pymc.io/t/var-names-not-working-with-sample-numpyro-nuts/10347)

<div class="topic-metadata">

**Author:** [@rtdew1](https://discourse.pymc.io/u/rtdew1)\
**Replies:** 5\
**Last updated:** [September 12, 2022, 8:21am UTC](https://discourse.pymc.io/t/var-names-not-working-with-sample-numpyro-nuts/10347 "2022-09-12T08:21:46Z")

</div>

When performing sampling with sample\_numpyro\_nuts on a model, and specifying the var\_names argument to selectively track variables, I get a mysterious error. MWE and error below. Any idea how to fix this? Or is it a bug? …

---

## [Pm.sampling\_jax to sample a MvNormal()](https://discourse.pymc.io/t/pm-sampling-jax-to-sample-a-mvnormal/10062)

<div class="topic-metadata">

**Author:** [@leotercas](https://discourse.pymc.io/u/leotercas)\
**Replies:** 4\
**Last updated:** [August 3, 2022, 10:58am UTC](https://discourse.pymc.io/t/pm-sampling-jax-to-sample-a-mvnormal/10062 "2022-08-03T10:58:06Z")

</div>

Hello! I have a code in PyMC3 using Theano Tensor and I’m trying to transcribe it to PyMC4 using Aesara. In my problem I want to sample the posterior distribution of a Normal Multivariate, in this case I get 50 observati…

[Next page](https://discourse.pymc.io/tag/jax/84.md?match_all_tags=true&page=1&tags%5B%5D=jax)