Replace with pt.split

Martin Ingram · Martin Ingram · commit cb070aa4d08d · 2025-08-29T18:38:55.000+02:00
diff --git a/pymc_extras/inference/dadvi/dadvi.py b/pymc_extras/inference/dadvi/dadvi.py
@@ -197,7 +197,8 @@ def create_dadvi_graph(
     )
 
     var_params = pt.vector(name="eta", shape=(2 * n_params,))
-    means, log_sds = var_params[:n_params], var_params[n_params:]
+
+    means, log_sds = pt.split(var_params, axis=0, splits_size=[n_params, n_params], n_splits=2)
 
     draw_matrix = pt.constant(draws)
     samples = means + pt.exp(log_sds) * draw_matrix

Original file line number	Diff line number	Diff line change
`@@ -197,7 +197,8 @@ def create_dadvi_graph(`
`197`	`197`	`)`
`198`	`198`
`199`	`199`	`var_params = pt.vector(name="eta", shape=(2 * n_params,))`
`200`		`- means, log_sds = var_params[:n_params], var_params[n_params:]`
	`200`	`+`
	`201`	`+ means, log_sds = pt.split(var_params, axis=0, splits_size=[n_params, n_params], n_splits=2)`
`201`	`202`
`202`	`203`	`draw_matrix = pt.constant(draws)`
`203`	`204`	`samples = means + pt.exp(log_sds) * draw_matrix`