More renaming

Author: fonnesbeck

pymc_extras/distributions/histogram_utils.py

@@ -15,6 +15,7 @@
 
 import numpy as np
 import pymc as pm
+import pymc_extras as pmx
 
 from numpy.typing import ArrayLike
 
@@ -115,7 +116,7 @@ def histogram_approximation(name, dist, *, observed, **h_kwargs):
     Discrete variables are reduced to unique repetitions (up to min_count)
 
     >>> import pymc as pm
-    >>> import pymc_experimental as pmx
+    >>> import pymc_extras as pmx
     >>> production = np.random.poisson([1, 2, 5], size=(1000, 3))
     >>> with pm.Model(coords=dict(plant=range(3))):
     ...     lam = pm.Exponential("lam", 1.0, dims="plant")

pymc_extras/model/marginal/distributions.py

@@ -16,7 +16,8 @@
 from pytensor.scan import scan
 from pytensor.tensor import TensorVariable
 
-from pymc_experimental.distributions import DiscreteMarkovChain
+from pymc_extras.distributions import DiscreteMarkovChain
+from pymc_extras.model.marginal.graph_analysis import get_support_axes
 
 
 class MarginalRV(OpFromGraph, MeasurableOp):
@@ -91,7 +92,7 @@ def reduce_batch_dependent_logps(
        as well as transpose the remaining axis of dep1 logp before adding the two element-wise.
 
     """
-    from pymc_experimental.model.marginal.graph_analysis import get_support_axes
+    from pymc_extras.model.marginal.graph_analysis import get_support_axes
 
     reduced_logps = []
     for dependent_op, dependent_logp, dependent_dims_connection in zip(

pymc_extras/model/marginal/graph_analysis.py

@@ -16,7 +16,7 @@
 from pytensor.tensor.subtensor import AdvancedSubtensor, Subtensor, get_idx_list
 from pytensor.tensor.type_other import NoneTypeT
 
-from pymc_experimental.model.marginal.distributions import MarginalRV
+from pymc_extras.model.marginal.distributions import MarginalRV
 
 
 def static_shape_ancestors(vars):

pymc_extras/statespace/filters/kalman_filter.py

@@ -11,12 +11,16 @@
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.slinalg import solve_triangular
 
-from pymc_experimental.statespace.filters.utilities import (
-    quad_form_sym,
-    split_vars_into_seq_and_nonseq,
+from pymc_extras.statespace.filters.utilities import (
+    compute_forecast_error_cov,
+    compute_kalman_gain,
+    compute_predicted_state,
+    compute_predicted_state_cov,
+    compute_updated_state,
+    compute_updated_state_cov,
     stabilize,
 )
-from pymc_experimental.statespace.utils.constants import JITTER_DEFAULT, MISSING_FILL
+from pymc_extras.statespace.utils.constants import JITTER_DEFAULT, MISSING_FILL
 
 MVN_CONST = pt.log(2 * pt.constant(np.pi, dtype="float64"))
 PARAM_NAMES = ["c", "d", "T", "Z", "R", "H", "Q"]

pymc_extras/statespace/filters/kalman_smoother.py

@@ -4,12 +4,12 @@
 from pytensor.compile import get_mode
 from pytensor.tensor.nlinalg import matrix_dot
 
-from pymc_experimental.statespace.filters.utilities import (
-    quad_form_sym,
-    split_vars_into_seq_and_nonseq,
+from pymc_extras.statespace.filters.utilities import (
+    compute_predicted_state,
+    compute_predicted_state_cov,
     stabilize,
 )
-from pymc_experimental.statespace.utils.constants import JITTER_DEFAULT
+from pymc_extras.statespace.utils.constants import JITTER_DEFAULT
 
 
 class KalmanSmoother:

pymc_extras/statespace/filters/utilities.py

@@ -2,11 +2,11 @@
 
 from pytensor.tensor.nlinalg import matrix_dot
 
-from pymc_experimental.statespace.core.representation import (
-    NEVER_TIME_VARYING,
-    VECTOR_VALUED,
+from pymc_extras.statespace.core.representation import (
+    PytensorRepresentation,
+    quad_form_sym,
 )
-from pymc_experimental.statespace.utils.constants import JITTER_DEFAULT
+from pymc_extras.statespace.utils.constants import JITTER_DEFAULT
 
 
 def decide_if_x_time_varies(x, name):

tests/__init__.py

@@ -11,3 +11,6 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+
+from pymc_extras.distributions import histogram_utils
+from pymc_extras.distributions.histogram_utils import histogram_approximation

tests/distributions/__init__.py

@@ -13,7 +13,7 @@
 #   limitations under the License.
 
 
-from pymc_experimental.distributions import histogram_utils
-from pymc_experimental.distributions.histogram_utils import histogram_approximation
+from pymc_extras.distributions import histogram_utils
+from pymc_extras.distributions.histogram_utils import histogram_approximation
 
 __all__ = ["histogram_utils", "histogram_approximation"]

tests/distributions/test_continuous.py

@@ -34,7 +34,7 @@
 )
 
 # the distributions to be tested
-from pymc_experimental.distributions import Chi, GenExtreme, Maxwell
+from pymc_extras.distributions import Chi, GenExtreme, Maxwell
 
 
 class TestGenExtremeClass:

tests/distributions/test_discrete.py

@@ -30,7 +30,7 @@
 )
 from pytensor import config
 
-from pymc_experimental.distributions import (
+from pymc_extras.distributions import (
     BetaNegativeBinomial,
     GeneralizedPoisson,
     Skellam,