fit_dadvi fails with gradient_backend="jax"

[jaj42]

Calling fit_dadvi(gradient_backend="jax") raises a numba.core.errors.TypingError during the posterior sampling step, even though the MAP optimization completes successfully.

Traceback (most recent call last):
  File "...", in fit_dadvi
    posterior, unconstrained_posterior = draws_from_laplace_approx(...)
  File "...", in draws_from_laplace_approx
    for out_draw in fn(mean, sigma):
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function numba_funcify_RandomVariable.<locals>.random ...>) found for signature:
 >>> random(readonly array(int64, 1d, C), NumPyRandomGeneratorType, readonly array(int64, 1d, C),
            readonly buffer(float64, 1d, C), array(float64, 1d, C))
...
  Rejected as the implementation raised a specific error:
    TypingError: Vectorized inputs must be arrays.

The reason is that the result of type jaxlib.ArrayImpl is propagated and not handled correctly by numba.

This can be fixed by the following patch:

diff --git a/pymc_extras/inference/dadvi/dadvi.py b/pymc_extras/inference/dadvi/dadvi.py
index b519499..1d620c6 100644
--- a/pymc_extras/inference/dadvi/dadvi.py
+++ b/pymc_extras/inference/dadvi/dadvi.py
@@ -185,7 +185,7 @@ def fit_dadvi(

     raveled_optimized = RaveledVars(result.x, dadvi_initial_point.point_map_info)

-    opt_var_params = result.x
+    opt_var_params = np.array(result.x)
     opt_means, opt_log_sds = np.split(opt_var_params, 2)

     posterior, unconstrained_posterior = draws_from_laplace_approx(

Environment

Python 3.13.7 (main, Sep  2 2025, 14:21:46) [Clang 20.1.4 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> pymc.__version__
'6.0.1'
>>> pymc_extras.__version__
'0.12.2.dev7+g71c1706c8'
>>> jax.__version__
'0.10.1'
>>> numba.__version__
'0.65.1'