rlouf / mcx

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.
https://rlouf.github.io/mcx
Apache License 2.0
325 stars 17 forks source link

Simplify the access to the graph's nodes #4

Open rlouf opened 4 years ago

rlouf commented 4 years ago

NetworkX's default is to separate the name of a node from the data attached to it; accessing a node by name G.nodes["n"] returns a dictionary with the data, here {"contents": Union[Var, Argument, RandVar, Transformation]}. As a result, the code is full of ["contents"] dictionary access calls, which is aesthetically unpleasant and unintuitive.

I think this behaviour can be overridden by subclassing the NodeView class in mcx.core.graphand overriding the graphs' nodes property.

This would greatly improve the code readability.