Add support for automatic naming of random variables

Author: zaxtax

pymc/distributions/distribution.py

@@ -22,6 +22,7 @@
 from typing import Callable, Optional, Sequence, Tuple, Union
 
 import numpy as np
+import opcode
 
 from aesara import tensor as at
 from aesara.compile.builders import OpFromGraph
@@ -164,6 +165,45 @@ def fn(*args, **kwargs):
     return fn
 
 
+# Helper function from pyprob
+def _extract_target_of_assignment(depth):
+    frame = sys._getframe(depth)
+    code = frame.f_code
+    next_instruction = code.co_code[frame.f_lasti + 2]
+    instruction_arg = code.co_code[frame.f_lasti + 3]
+    instruction_name = opcode.opname[next_instruction]
+    if instruction_name == "STORE_FAST":
+        return code.co_varnames[instruction_arg]
+    elif instruction_name in ["STORE_NAME", "STORE_GLOBAL"]:
+        return code.co_names[instruction_arg]
+    elif (
+        instruction_name in ["LOAD_FAST", "LOAD_NAME", "LOAD_GLOBAL"]
+        and opcode.opname[code.co_code[frame.f_lasti + 4]] in ["LOAD_CONST", "LOAD_FAST"]
+        and opcode.opname[code.co_code[frame.f_lasti + 6]] == "STORE_SUBSCR"
+    ):
+        if instruction_name == "LOAD_FAST":
+            base_name = code.co_varnames[instruction_arg]
+        else:
+            base_name = code.co_names[instruction_arg]
+
+        second_instruction = opcode.opname[code.co_code[frame.f_lasti + 4]]
+        second_arg = code.co_code[frame.f_lasti + 5]
+        if second_instruction == "LOAD_CONST":
+            value = code.co_consts[second_arg]
+        elif second_instruction == "LOAD_FAST":
+            var_name = code.co_varnames[second_arg]
+            value = frame.f_locals[var_name]
+        else:
+            value = None
+        if value is not None:
+            index_name = repr(value)
+            return base_name + "[" + index_name + "]"
+        else:
+            return None
+    else:
+        return None
+
+
 class SymbolicRandomVariable(OpFromGraph):
     """Symbolic Random Variable
 
@@ -216,7 +256,6 @@ class Distribution(metaclass=DistributionMeta):
 
     def __new__(
         cls,
-        name: str,
         *args,
         rng=None,
         dims: Optional[Dims] = None,
@@ -234,8 +273,6 @@ def __new__(
         ----------
         cls : type
             A PyMC distribution.
-        name : str
-            Name for the new model variable.
         rng : optional
             Random number generator to use with the RandomVariable.
         dims : tuple, optional
@@ -277,6 +314,19 @@ def __new__(
                 "for a standalone distribution."
             )
 
+        if "name" in kwargs:
+            name = kwargs.pop("name")
+        elif len(args) > 0 and isinstance(args[0], string_types):
+            name = args[0]
+            args = args[1:]
+        else:
+            name = _extract_target_of_assignment(2)
+            if name is None:
+                raise TypeError("Name could not be inferred for variable")
+
+        if not isinstance(name, string_types):
+            raise TypeError(f"Name needs to be a string but got: {name}")
+
         if "testval" in kwargs:
             initval = kwargs.pop("testval")
             warnings.warn(
@@ -285,9 +335,6 @@ def __new__(
                 stacklevel=2,
             )
 
-        if not isinstance(name, string_types):
-            raise TypeError(f"Name needs to be a string but got: {name}")
-
         dims = convert_dims(dims)
         if observed is not None:
             observed = convert_observed_data(observed)

pymc/tests/distributions/test_distribution.py

@@ -416,3 +416,23 @@ def test_tag_future_warning_dist():
         with pytest.warns(FutureWarning, match="Use model.rvs_to_values"):
             value_var = new_x.tag.value_var
         assert value_var == "1"
+
+
+def test_autonaming():
+    """Test that random variable ends up with same name as assignment variable"""
+    d = {}
+    with pm.Model() as m:
+        x = pm.Normal(0.0, 1.0)
+        d[2] = pm.Normal(0.0, 1.0)
+
+    assert x.name == "x"
+    assert m["x"] == x
+    assert d[2].name == "d[2]"
+    assert m["d[2]"] == d[2]
+
+
+def test_autonaming_noname():
+    """Test that autonaming fails if no assignment can be found"""
+    with pytest.raises(TypeError, match="Name could not be inferred for variable"):
+        with pm.Model():
+            pm.Normal()