avoid jaxifying logp multiple times

nataziel · ricardoV94 · commit 2fbc8a927fce · 2025-02-26T07:00:44.000+01:00
diff --git a/pymc/sampling/jax.py b/pymc/sampling/jax.py
@@ -669,6 +669,7 @@ def sample_jax_nuts(
         random_seed=random_seed,
         initial_points=initial_points,
         nuts_kwargs=nuts_kwargs,
+        logp_fn=logp_fn,
     )
     tic2 = datetime.now()
 

Original file line number	Diff line number	Diff line change
`@@ -669,6 +669,7 @@ def sample_jax_nuts(`
`669`	`669`	`random_seed=random_seed,`
`670`	`670`	`initial_points=initial_points,`
`671`	`671`	`nuts_kwargs=nuts_kwargs,`
	`672`	`+ logp_fn=logp_fn,`
`672`	`673`	`)`
`673`	`674`	`tic2 = datetime.now()`
`674`	`675`