Add experimental dims module with objects that follow dim-based semantics (like xarray without coordinates)

.github/workflows/tests.yml

@@ -136,6 +136,12 @@ jobs:
             tests/logprob/test_transforms.py
             tests/logprob/test_utils.py
 
+          - |
+            tests/dims/distributions/test_core.py
+            tests/dims/distributions/test_scalar.py
+            tests/dims/distributions/test_vector.py
+            tests/dims/test_model.py
+
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:

conda-envs/environment-alternative-backends.yml

@@ -22,7 +22,7 @@ dependencies:
 - numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.7,<2.32
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/environment-dev.yml

@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.7,<2.32
 - python-graphviz
 - networkx
 - scipy>=1.4.1

conda-envs/environment-docs.yml

@@ -11,7 +11,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.7,<2.32
 - python-graphviz
 - rich>=13.7.1
 - scipy>=1.4.1

conda-envs/environment-test.yml

@@ -14,7 +14,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.7,<2.32
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/windows-environment-dev.yml

@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.7,<2.32
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/windows-environment-test.yml

@@ -15,7 +15,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=2.31.2,<2.32
+- pytensor>=2.31.7,<2.32
 - python-graphviz
 - networkx
 - rich>=13.7.1

docs/source/api.rst

@@ -23,6 +23,7 @@ API
    api/backends
    api/misc
    api/testing
+   api/dims
 
 ------------------
 Dimensionality

docs/source/api/dims.rst

@@ -0,0 +1,18 @@
+.. _api_dims:
+
+Dims
+====
+
+.. warning:: This module is experimental and may contain critical breaks. API changes are expected in future releases.
+
+This submodule contains functions for defining distributions and operations that use explicit dimensions.
+
+The module is presented in :ref:`dims_module`.
+
+
+.. toctree::
+
+   dims/model
+   dims/math
+   dims/distributions
+   dims/transforms

docs/source/api/dims/distributions.rst

@@ -0,0 +1,38 @@
+*************
+Distributions
+*************
+
+Scalar distributions
+====================
+
+.. currentmodule:: pymc.dims
+.. autosummary::
+   :toctree: generated/
+   :template: distribution.rst
+
+   Flat
+   HalfFlat
+   Normal
+   HalfNormal
+   LogNormal
+   StudentT
+   HalfStudentT
+   Cauchy
+   HalfCauchy
+   Beta
+   Laplace
+   Gamma
+   InverseGamma
+
+
+Vector distributions
+====================
+
+.. currentmodule:: pymc.dims
+.. autosummary::
+   :toctree: generated/
+   :template: distribution.rst
+
+   Categorical
+   MvNormal
+   ZeroSumNormal

docs/source/api/dims/math.rst

@@ -0,0 +1,9 @@
+***************************************
+Mathematical operations with dimensions
+***************************************
+
+This module wraps all the mathematical operations defined in :ref:`pytensor.xtensor.math <pytensor:libdoc_xtensor_math>`.
+
+It includes a ``linalg`` submodule that wraps all the operations defined in :ref:`pytensor.xtensor.linalg <pytensor:libdoc_xtensor_linalg>`.
+
+Operations defined at the module level in :ref:`pytensor.xtensor <pytensor:libdoc_xtensor_module_function>` are available at the ``pymc.dims`` module level.

docs/source/api/dims/model.rst

@@ -0,0 +1,11 @@
+******************
+Model constructors
+******************
+
+.. currentmodule:: pymc.dims
+.. autosummary::
+   :toctree: generated/
+
+   Data
+   Deterministic
+   Potential

docs/source/api/dims/transforms.rst

@@ -0,0 +1,11 @@
+***********************
+Distribution Transforms
+***********************
+
+.. currentmodule:: pymc.dims.transforms
+.. autosummary::
+   :toctree: generated/
+
+   LogTransform
+   LogOddsTransform
+   ZeroSumTransform