Of course @fritzo, I could explain the process I have gone through.

First, I started by studying the unit tests and seeing what the get_model_relations function was expected to return depending on the model, which returned a dictionary representing a graph. Next, I started debugging inside that method the logic you had implemented and started to understand how the different parts "exchange messages" (dictionaries) and I noticed that everything resided inside the apply_stack call.

When calling _post_process_message(), it checks if the TrackProvenance object has a method f"_pyro_post_{}".format(msg["type"]), but as being deterministic, it was not implemented yet, so I need to implement it. To do it, I just remembered watching the same structured code somewhere in Numpyro's code, so that's what I applied. In the same way, the _pyro_post_deterministic of TraceMessenger was implemented.

Of course, the primitive deterministic should be modified to match the above. I just took another look to Numpyro's deterministic to take some ideas. I noticed that their implementation differs, so I gave a try and implemented something similar. I ended up with the idea submitted, emulating the same functionality.

Finally, I apply some corrections to the core of get_model_relations to recognize deterministic node's parents and ended applying the same visualization (dashed) to the rendered graph as Numpyro.

That`s it. Hope it works, or at least, it do the job on my project 🤓

About the merge intention, I just take the render-deterministic branch to continue the work done. It cannot be merged into the dev branch, can it? Or can it?

Hi, any advances on this?

@r3v1 could you revert your changes to the body of pyro.deterministic and refactor your new handler methods for deterministic to fit within handler methods for sample? Changing the type of deterministic will break some things elsewhere in Pyro and downstream.

You can just use a helper function to check within a _pyro_sample/_pyro_post_sample if a sample site came from deterministic:

def site_is_deterministic(msg: dict) -> bool:
    return msg["type"] == "sample" and msg["infer"].get("_deterministic", False)

...

class ...:
    ...
    def _pyro_sample(self, msg):
        if site_is_deterministic(msg):
            # new logic for handling deterministic sites here
            ...
            return

        # old logic for handling sample sites here
        ...

@r3v1 I also just changed the base branch to dev, can you resolve the merge conflicts and push so that CI can run tests?

Ok, changes made in pull request #3266 to dev branch, so this branch may be closed