[DO NOT MERGE WITHOUT SQUASHING with Tweaks to SymbolicRandomVariables] Get rid of magic

Author: ricardoV94

pymc/distributions/distribution.py

@@ -272,21 +272,6 @@ def fn(*args, **kwargs):
     return fn
 
 
-class _class_or_instancemethod(classmethod):
-    """Allow a method to be called both as a classmethod and an instancemethod,
-    giving priority to the instancemethod.
-
-    This is used to allow extracting information from the signature of a SymbolicRandomVariable
-    which may be provided either as a class attribute or as an instance attribute.
-
-    Adapted from https://stackoverflow.com/a/28238047
-    """
-
-    def __get__(self, instance, type_):
-        descr_get = super().__get__ if instance is None else self.__func__.__get__
-        return descr_get(instance, type_)
-
-
 class SymbolicRandomVariable(OpFromGraph):
     """Symbolic Random Variable
 
@@ -343,37 +328,40 @@ def _parse_params_signature(signature):
         else:
             return _parse_gufunc_signature(params_signature)
 
-    @_class_or_instancemethod
-    @property
-    def ndims_params(cls_or_self) -> Sequence[int] | None:
+    @staticmethod
+    def _ndims_params_from_signature(signature: str) -> Sequence[int]:
         """Number of core dimensions of the distribution's parameters."""
-        signature = cls_or_self.signature
-        if signature is None:
-            return None
-        inputs_signature, _ = cls_or_self._parse_params_signature(signature)
+        inputs_signature, _ = SymbolicRandomVariable._parse_params_signature(signature)
         return [len(sig) for sig in inputs_signature]
 
-    @_class_or_instancemethod
+    @classmethod
     @property
-    def ndim_supp(cls_or_self) -> int | None:
-        """Number of support dimensions of the RandomVariable
-
-        (0 for scalar, 1 for vector, ...)
-        """
-        signature = cls_or_self.signature
+    def ndims_params(cls) -> Sequence[int] | None:
+        signature = cls.signature
         if signature is None:
             return None
-        _, outputs_params_signature = cls_or_self._parse_params_signature(signature)
+        return cls._ndims_params_from_signature(signature)
+
+    @staticmethod
+    def _ndim_supp_from_signature(signature: str) -> int:
+        _, outputs_params_signature = SymbolicRandomVariable._parse_params_signature(signature)
         return max(len(out_sig) for out_sig in outputs_params_signature)
 
-    @_class_or_instancemethod
+    @classmethod
     @property
-    def default_output(cls_or_self) -> int | None:
-        signature = cls_or_self.signature
+    def ndim_supp(cls) -> int | None:
+        """Number of support dimensions of the RandomVariable
+
+        (0 for scalar, 1 for vector, ...)
+        """
+        signature = cls.signature
         if signature is None:
             return None
+        return cls._ndim_supp_from_signature(signature)
 
-        _, outputs_signature = cls_or_self._parse_signature(signature)
+    @staticmethod
+    def _default_output_from_signature(signature: str) -> int | None:
+        _, outputs_signature = SymbolicRandomVariable._parse_signature(signature)
 
         # If there is a single non `[rng]` outputs, that is the default one!
         candidate_default_output = [
@@ -384,6 +372,14 @@ def default_output(cls_or_self) -> int | None:
         else:
             return None
 
+    @classmethod
+    @property
+    def default_output(cls) -> int | None:
+        signature = cls.signature
+        if signature is None:
+            return None
+        return cls._default_output_from_signature(signature)
+
     @staticmethod
     def get_idxs(signature: str) -> tuple[tuple[int], int | None, tuple[int]]:
         """Parse signature and return indexes for *[rng], [size] and parameters"""
@@ -406,10 +402,17 @@ def __init__(
         **kwargs,
     ):
         """Initialize a SymbolicRandomVariable class."""
-        if "signature" in kwargs:
-            self.signature = kwargs.pop("signature")
-
-        if "ndim_supp" in kwargs:
+        signature = kwargs.pop("signature", None)
+        if signature is not None:
+            self.signature = signature
+            # Override class properties with instance properties
+            self.ndims_params = self._ndims_params_from_signature(signature)
+            self.ndim_supp = self._ndim_supp_from_signature(self.signature)
+            self.default_output = self._default_output_from_signature(self.signature)
+
+        elif "ndim_supp" in kwargs:
+            # For backwards compatibility we allow passing ndim_supp without signature
+            # This is the only variable that PyMC absolutely needs to work with SymbolicRandomVariables
             self.ndim_supp = kwargs.pop("ndim_supp")
 
         if self.ndim_supp is None: