Respond to feedback

jessegrabowski · jessegrabowski · commit 427dfb85b938 · 2025-09-23T07:51:37.000-05:00
diff --git a/pymc_extras/inference/dadvi/dadvi.py b/pymc_extras/inference/dadvi/dadvi.py
@@ -341,7 +341,7 @@ def dadvi_result_to_idata(
     all_varnames = [
         x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=True)
     ]
-    unconstrained_names = set(all_varnames) - set(constrained_names)
+    unconstrained_names = sorted(set(all_varnames) - set(constrained_names))
 
     idata = az.InferenceData(posterior=transformed_result[constrained_names])
 
diff --git a/pymc_extras/inference/laplace_approx/idata.py b/pymc_extras/inference/laplace_approx/idata.py
@@ -115,7 +115,7 @@ def map_results_to_inference_data(
         x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=True)
     ]
 
-    unconstrained_names = set(all_varnames) - set(constrained_names)
+    unconstrained_names = sorted(set(all_varnames) - set(constrained_names))
 
     idata = az.from_dict(
         posterior={
diff --git a/tests/inference/dadvi/test_dadvi.py b/tests/inference/dadvi/test_dadvi.py
@@ -15,7 +15,7 @@ def rng():
 
 @pytest.mark.parametrize(
     "mode, gradient_backend",
-    [(None, "pytensor"), ("NUMBA", "pytensor"), ("JAX", "jax"), ("JAX", "pytensor")],
+    [(None, "pytensor")],
 )
 def test_fit_dadvi_basic(mode, gradient_backend):
     # Example from BDA3, section 4.1 (same as Laplace test)

Original file line number	Diff line number	Diff line change
`@@ -341,7 +341,7 @@ def dadvi_result_to_idata(`
`341`	`341`	`all_varnames = [`
`342`	`342`	`x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=True)`
`343`	`343`	`]`
`344`		`- unconstrained_names = set(all_varnames) - set(constrained_names)`
	`344`	`+ unconstrained_names = sorted(set(all_varnames) - set(constrained_names))`
`345`	`345`
`346`	`346`	`idata = az.InferenceData(posterior=transformed_result[constrained_names])`
`347`	`347`
Original file line number	Diff line number	Diff line change
`@@ -115,7 +115,7 @@ def map_results_to_inference_data(`
`115`	`115`	`x.name for x in get_default_varnames(model.unobserved_value_vars, include_transformed=True)`
`116`	`116`	`]`
`117`	`117`
`118`		`- unconstrained_names = set(all_varnames) - set(constrained_names)`
	`118`	`+ unconstrained_names = sorted(set(all_varnames) - set(constrained_names))`
`119`	`119`
`120`	`120`	`idata = az.from_dict(`
`121`	`121`	`posterior={`
Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ def rng():`
`15`	`15`
`16`	`16`	`@pytest.mark.parametrize(`
`17`	`17`	`"mode, gradient_backend",`
`18`		`- [(None, "pytensor"), ("NUMBA", "pytensor"), ("JAX", "jax"), ("JAX", "pytensor")],`
	`18`	`+ [(None, "pytensor")],`
`19`	`19`	`)`
`20`	`20`	`def test_fit_dadvi_basic(mode, gradient_backend):`
`21`	`21`	`# Example from BDA3, section 4.1 (same as Laplace test)`