Fix bug in transform_scan_values

Author: ricardoV94

pymc/logprob/transforms.py

@@ -237,7 +237,7 @@ def transform_scan_values(fgraph: FunctionGraph, node: Node) -> Optional[List[No
         return None
 
     transforms = [
-        values_to_transforms.get(rv_map_feature.original_values[value], None)
+        values_to_transforms.get(rv_map_feature.original_values[value_var], None)
         for value_var in value_vars
     ]
 

tests/logprob/test_transforms.py

@@ -914,8 +914,13 @@ def test_scan_transform():
     init = at.random.beta(1, 1, name="init")
     init_vv = init.clone()
 
+    def scan_step(prev_innov):
+        next_innov = at.random.beta(prev_innov * 10, (1 - prev_innov) * 10)
+        update = {next_innov.owner.inputs[0]: next_innov.owner.outputs[0]}
+        return next_innov, update
+
     innov, _ = scan(
-        fn=lambda prev_innov: at.random.beta(prev_innov * 10, (1 - prev_innov) * 10),
+        fn=scan_step,
         outputs_info=[init],
         n_steps=4,
     )