Add pmd.CustomDist — dims-aware custom distribution for pymc.dims

.github/workflows/tests.yml

@@ -142,6 +142,7 @@ jobs:
           - |
             tests/dims/distributions/test_core.py
             tests/dims/distributions/test_censored.py
+            tests/dims/distributions/test_custom.py
             tests/dims/distributions/test_scalar.py
             tests/dims/distributions/test_vector.py
             tests/dims/test_model.py

pymc/dims/distributions/__init__.py

@@ -12,5 +12,6 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 from pymc.dims.distributions.censored import Censored
+from pymc.dims.distributions.custom import CustomDist
 from pymc.dims.distributions.scalar import *
 from pymc.dims.distributions.vector import *

pymc/dims/distributions/core.py

@@ -28,7 +28,7 @@
 from pytensor.xtensor.basic import XTensorFromTensor, xtensor_from_tensor
 from pytensor.xtensor.shape import Transpose
 from pytensor.xtensor.type import XTensorVariable
-from pytensor.xtensor.vectorization import XRV
+from pytensor.xtensor.vectorization import XRV, XElemwise
 
 from pymc import SymbolicRandomVariable, modelcontext
 from pymc.dims.distributions.transforms import DimTransform, log_odds_transform, log_transform
@@ -374,6 +374,14 @@ def expand_dist_dims(dist: XTensorVariable, extra_dims: dict[str, Any]) -> XTens
             # We don't propagate the old RNG, because we don't want the new and old dists to be correlated
             new_rng = pt.random.shared_rng(seed=None)
             return new_dist_op(new_rng, *extra_dims.values(), *params_and_dim_lengths)
+        case XElemwise():
+            expanded_inputs = [
+                expand_dist_dims(inp, extra_dims=extra_dims)
+                if isinstance(inp, XTensorVariable)
+                else inp
+                for inp in dist.owner.inputs
+            ]
+            return dist.owner.op.make_node(*expanded_inputs).outputs[0]
         case Transpose():
             return expand_dist_dims(dist.owner.inputs[0], extra_dims=extra_dims).transpose(
                 ..., *dist.dims