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()