rename pymc_experimental -> pymc_extras

jessegrabowski · jessegrabowski · commit 03593734c064 · 2025-01-02T14:07:24.000+08:00
diff --git a/pymc_extras/model/marginal/distributions.py b/pymc_extras/model/marginal/distributions.py
@@ -7,7 +7,6 @@
 from pymc.logprob.abstract import MeasurableOp, _logprob
 from pymc.logprob.basic import conditional_logp, logp
 from pymc.pytensorf import constant_fold
-from pytensor import Variable
 from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.mode import Mode
 from pytensor.graph import Op, vectorize_graph
diff --git a/pymc_extras/sampling/__init__.py b/pymc_extras/sampling/__init__.py
diff --git a/pymc_extras/sampling/mcmc.py b/pymc_extras/sampling/mcmc.py
@@ -4,7 +4,7 @@
 from pymc.sampling.mcmc import sample
 from pytensor.graph.rewriting.basic import GraphRewriter
 
-from pymc_experimental.sampling.optimizations.optimize import (
+from pymc_extras.sampling.optimizations.optimize import (
     TAGS_TYPE,
     optimize_model_for_mcmc_sampling,
 )
diff --git a/pymc_extras/sampling/optimizations/__init__.py b/pymc_extras/sampling/optimizations/__init__.py
@@ -1,9 +1,8 @@
 # ruff: noqa: F401
 # Add rewrites to the optimization DBs
-import pymc_experimental.sampling.optimizations.conjugacy
-import pymc_experimental.sampling.optimizations.summary_stats
 
-from pymc_experimental.sampling.optimizations.optimize import (
+from pymc_extras.sampling.optimizations import conjugacy, summary_stats
+from pymc_extras.sampling.optimizations.optimize import (
     optimize_model_for_mcmc_sampling,
     posterior_optimization_db,
 )
diff --git a/pymc_extras/sampling/optimizations/conjugacy.py b/pymc_extras/sampling/optimizations/conjugacy.py
@@ -10,10 +10,10 @@
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.subtensor import _sum_grad_over_bcasted_dims as sum_bcasted_dims
 
-from pymc_experimental.sampling.optimizations.conjugate_sampler import (
+from pymc_extras.sampling.optimizations.conjugate_sampler import (
     ConjugateRV,
 )
-from pymc_experimental.sampling.optimizations.optimize import posterior_optimization_db
+from pymc_extras.sampling.optimizations.optimize import posterior_optimization_db
 
 
 def register_conjugacy_rewrites_variants(rewrite_fn, tracks=(ModelFreeRV,)):
diff --git a/pymc_extras/sampling/optimizations/conjugate_sampler.py b/pymc_extras/sampling/optimizations/conjugate_sampler.py
@@ -13,7 +13,7 @@
 from pytensor.link.jax.linker import JAXLinker
 from pytensor.tensor.random.type import RandomGeneratorType
 
-from pymc_experimental.utils.ofg import inline_ofg_outputs
+from pymc_extras.utils.ofg import inline_ofg_outputs
 
 
 class ConjugateRV(OpFromGraph, MeasurableOp):
diff --git a/pymc_extras/sampling/optimizations/optimize.py b/pymc_extras/sampling/optimizations/optimize.py
diff --git a/pymc_extras/sampling/optimizations/summary_stats.py b/pymc_extras/sampling/optimizations/summary_stats.py
@@ -5,7 +5,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import node_rewriter
 
-from pymc_experimental.sampling.optimizations.optimize import posterior_optimization_db
+from pymc_extras.sampling.optimizations.optimize import posterior_optimization_db
 
 
 @node_rewriter(tracks=[ModelObservedRV])
diff --git a/pymc_extras/utils/ofg.py b/pymc_extras/utils/ofg.py
@@ -5,7 +5,7 @@
 from pytensor.graph.replace import clone_replace
 
 
-def inline_ofg_outputs(op: OpFromGraph, inputs: Sequence[Variable]) -> tuple[Variable]:
+def inline_ofg_outputs(op: OpFromGraph, inputs: Sequence[Variable]) -> list[Variable]:
     """Inline the inner graph (outputs) of an OpFromGraph Op.
 
     Whereas `OpFromGraph` "wraps" a graph inside a single Op, this function "unwraps"
diff --git a/tests/sampling/mcmc/test_mcmc.py b/tests/sampling/mcmc/test_mcmc.py
@@ -8,7 +8,7 @@
 from pymc.sampling.mcmc import sample
 from pymc.step_methods import Slice
 
-from pymc_experimental import opt_sample
+from pymc_extras import opt_sample
 
 
 def test_custom_step_raises():
diff --git a/tests/sampling/optimizations/test_conjugacy.py b/tests/sampling/optimizations/test_conjugacy.py
@@ -6,8 +6,8 @@
 from pymc.model.transform.conditioning import remove_value_transforms
 from pymc.sampling import draw
 
-from pymc_experimental.sampling.optimizations.conjugate_sampler import ConjugateRV
-from pymc_experimental.sampling.optimizations.optimize import optimize_model_for_mcmc_sampling
+from pymc_extras.sampling.optimizations import optimize_model_for_mcmc_sampling
+from pymc_extras.sampling.optimizations.conjugate_sampler import ConjugateRV
 
 
 @pytest.mark.parametrize("eager", [False, True])
diff --git a/tests/sampling/optimizations/test_summary_stats.py b/tests/sampling/optimizations/test_summary_stats.py
@@ -3,7 +3,7 @@
 from pymc.distributions import HalfNormal, Normal
 from pymc.model.core import Model
 
-from pymc_experimental.sampling.optimizations.optimize import optimize_model_for_mcmc_sampling
+from pymc_extras.sampling.optimizations import optimize_model_for_mcmc_sampling
 
 
 def test_summary_stats_normal():

Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@`
`4`	`4`	`from pymc.sampling.mcmc import sample`
`5`	`5`	`from pytensor.graph.rewriting.basic import GraphRewriter`
`6`	`6`
`7`		`-from pymc_experimental.sampling.optimizations.optimize import (`
	`7`	`+from pymc_extras.sampling.optimizations.optimize import (`
`8`	`8`	`TAGS_TYPE,`
`9`	`9`	`optimize_model_for_mcmc_sampling,`
`10`	`10`	`)`