Surprising edge categorization if there's a hidden intermediate calculator

[jobrachem]

Consider the following model code and corresponding graph:

import liesel.model as lsl
import tensorflow_probability.substrates.jax.distributions as tfd
import jax.numpy as jnp

rent = jnp.linspace(-2, 2, 50)

mu = lsl.param(0.0, name="mu")
sigma_sq = lsl.param(1.0, name="sigma_sq")
sigma_calc = lsl.Calc(jnp.sqrt, sigma_sq)

# sigma = lsl.Var(sigma_calc, name="sigma")

y_dist = lsl.Dist(tfd.Normal, loc=mu, scale=sigma_calc)
y = lsl.obs(rent, distribution=y_dist, name="rent")

gb = lsl.GraphBuilder().add(y)
gb.plot_vars()

I am surprised by the fact the the edge from sigma_sq to rent is colored black, i.e. it is categorized as "used in value" instead of "used in distribution". I would expect it to be colored grey.

If we include sigma as a lsl.Var, we can see what is going on. sigma_sq is used in the value of sigma. And then sigma is used in the distribution of rent.

So the surprising behavior stems from how the edge categorization deals with nodes that are hidden in the variables plot. @GianmarcoCallegher can you think of a way to solve this?

[jobrachem]

See #89

[GianmarcoCallegher]

@jobrachem it is ready for the review. I didn't adopt the solution proposed in #89 (probably because it doesn't work). This is a test file:

tests_file.py.zip