rename pymc_experimental -> pymc_extras

Author: jessegrabowski

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

pymc_experimental/sampling/__init__.py renamed to pymc_extras/sampling/__init__.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,
 )

pymc_experimental/sampling/mcmc.py renamed to pymc_extras/sampling/mcmc.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,
 )

pymc_experimental/sampling/optimizations/__init__.py renamed to pymc_extras/sampling/optimizations/__init__.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,)):

pymc_experimental/sampling/optimizations/conjugacy.py renamed to pymc_extras/sampling/optimizations/conjugacy.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):

pymc_experimental/sampling/optimizations/conjugate_sampler.py renamed to pymc_extras/sampling/optimizations/conjugate_sampler.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])

pymc_experimental/sampling/optimizations/optimize.py renamed to pymc_extras/sampling/optimizations/optimize.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"

pymc_experimental/sampling/optimizations/summary_stats.py renamed to pymc_extras/sampling/optimizations/summary_stats.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():