Topics tagged jax

Topic	Replies	Views	Activity
Possible bug: sample_prior_predictive with JAX fails for time-varying statespace models General bug , jax , state_space	2	111	February 17, 2026
Integrating custom JAX-based likelihoods (e.g. auto differentiable stochastic filters) w/ PyMC sampling workflow General jax , ode_pde	2	148	September 22, 2025
Pymc State Space .1.4 -- XlaRuntimeError Incorrect output dtype for return value #0: Expected: int64, Actual: int32 version agnostic time_series , jax , state_space	10	500	February 27, 2025
AttributeError: 'numpy.ndarray' object has no attribute 'at' when sampling LKJ Cholesky Covariance Priors for Multivariate Normal Models example with numpyro or blackjax v5 gpu , bug , jax , pytensor	4	315	February 26, 2025
Numpyro sampler issue version agnostic jax	1	135	February 5, 2025
How to write a PyTensor Op to wrap Jax ODEs with multiple input parameters v5 jax	7	1038	July 28, 2024
Jax-Metal pm.Sample with numpyro on M1/M2 v5 jax	1	604	June 12, 2024
New distribution for PyMC: alpha-stable family [beta] Sharing modeling , jax	5	335	May 28, 2024
Recent error from Jax's (numpy) clip method v5 modeling , jax	2	447	May 21, 2024
Jax sampling for bayesian neural network v5 jax , pytensor	3	331	March 27, 2024
Can jax be used to wrap a function that contained python libraries? v5 jax	2	398	December 25, 2023
How to resolve Input Dimension Mis-match Error in Hierarchical Bayesian Inference with PyMC3 v5 theano , modeling , jax , hierarchical	3	427	December 20, 2023
pm.Categorical with sample_numpyro_nuts v5 jax	3	439	November 24, 2023
Dynamic shaping, "round" function, JAX, and a "few" more questions v5 modeling , jax , pytensor	28	1937	September 29, 2023
JAX Sampling Error with TruncatedNormal Distribution Questions jax	6	1077	July 27, 2023
Numpyro JAX sampling very slow v5 development , theano , bug , modeling , jax	1	1843	July 9, 2023
Jax on Apple Silicon GPU version agnostic jax	2	2312	June 13, 2023
NameError: unbound axis name raised during transformation of variables after sample_numpyro_nuts v5 bug , jax	6	1087	June 13, 2023
Wild results for masked multinomial with jax sampler(s) v5 modeling , jax	2	430	May 7, 2023
Pymc3 on GPU using jax v3 jax	2	1659	April 20, 2023
Out of Memory when using pm.sampling.jax.sample_blackjax_nuts v5 gpu , modeling , jax	2	979	March 23, 2023
How can I output a gradient in vector format in Op.grad instance? v5 jax	7	964	January 14, 2023
How to use JAX ODEs and Neural Networks in PyMC Sharing jax	0	450	January 4, 2023
Numpyro Convergence Diagnostics v5 jax	1	754	December 30, 2022
Out of memory when "transforming variables" in Numpyro & JAX v5 jax	10	1451	December 12, 2022
Sample_numpyro_nuts hangs when parallelizing over datasets with multiprocessing v5 jax	0	734	November 9, 2022
JAX Sampling Error when same model had previously worked without it v5 jax	1	413	October 25, 2022
Slow inference for numpyro sampling on Colab GPU v5 jax	3	939	October 5, 2022
Var_names not working with sample_numpyro_nuts v5 bug , jax	5	660	September 12, 2022
Pm.sampling_jax to sample a MvNormal() v5 modeling , sampling , jax	4	774	August 3, 2022