add desuffix to trace to extract forward/reverse variables

coix/core.py

@@ -117,6 +117,8 @@ def desuffix(trace):
   new_trace = {}
   for name in trace:
     raw_name = names_to_raw_names[name]
+    if raw_name != name and isinstance(trace[name], dict):
+      trace[name]["desuffix"] = True
     new_trace[name[: len(name) - num_suffix_min[raw_name]]] = trace[name]
   return new_trace
 

coix/loss.py

@@ -124,7 +124,7 @@ def fkl_loss(q_trace, p_trace, incoming_log_weight, incremental_log_weight):
   proposal_sites = [
       name
       for name, site in q_trace.items()
-      if util.is_observed_site(site) or name.endswith("_PREV_")
+      if isinstance(site, dict) and "desuffix" in site
   ]
 
   proposal_lp = sum(
@@ -171,7 +171,7 @@ def rkl_loss(q_trace, p_trace, incoming_log_weight, incremental_log_weight):
   proposal_sites = [
       name
       for name, site in q_trace.items()
-      if util.is_observed_site(site) or name.endswith("_PREV_")
+      if isinstance(site, dict) and "desuffix" in site
   ]
 
   proposal_lp = sum(
@@ -208,7 +208,7 @@ def rws_loss(q_trace, p_trace, incoming_log_weight, incremental_log_weight):
   proposal_sites = [
       name
       for name, site in q_trace.items()
-      if util.is_observed_site(site) or name.endswith("_PREV_")
+      if isinstance(site, dict) and "desuffix" in site
   ]
 
   proposal_lp = sum(

coix/util.py

@@ -223,32 +223,6 @@ def step_fn(params, opt_state, *args, **kwargs):
   return params, metrics
 
 
-def _remove_suffix(name):
-  i = 0
-  while name.endswith("_PREV_"):
-    i += len("_PREV_")
-    name = name[: -len("_PREV_")]
-  return name, i
-
-
-def desuffix(trace):
-  """Remove unnecessary suffix terms added to the trace."""
-  names_to_raw_names = {}
-  num_suffix_min = {}
-  for name in trace:
-    raw_name, num_suffix = _remove_suffix(name)
-    names_to_raw_names[name] = raw_name
-    if raw_name in num_suffix_min:
-      num_suffix_min[raw_name] = min(num_suffix_min[raw_name], num_suffix)
-    else:
-      num_suffix_min[raw_name] = num_suffix
-  new_trace = {}
-  for name in trace:
-    raw_name = names_to_raw_names[name]
-    new_trace[name[: len(name) - num_suffix_min[raw_name]]] = trace[name]
-  return new_trace
-
-
 def get_batch_ndims(xs):
   """Gets the number of same-size leading dimensions of the elements in xs."""
   if not xs: