Inspect a model’s graph

A module is a flat graph of typed bindings. This page lists the bindings with what inference knows about them, follows dependencies up and down, cuts the model down to what one output needs, renders it as mathematics, and counts the operations of a compiled query.

Runtime budget: 60 s.

from flatppl import flatppl

model = flatppl("""
mu = elementof(reals)
tau = elementof(posreals)
theta = elementof(cartpow(reals, 4))
scaled = mu .+ tau .* theta
total = sum(scaled)
spread = sum(theta .* theta)
""")

Each binding reports its kind, type, shape, domain and source position.

for binding in model.bindings.values():
    print(binding)
print(model.parameters)

<Binding mu: real parameter, line 2: elementof(reals)>
<Binding tau: real parameter, line 3: elementof(posreals)>
<Binding theta: real[4] parameter, line 4: elementof(cartpow(reals, 4))>
<Binding scaled: real[4] derived, line 5: mu .+ tau .* theta>
<Binding total: real derived, line 6: sum(scaled)>
<Binding spread: real derived, line 7: sum(theta .* theta)>
<Parameters mu, tau, theta>

Dependencies, one step at a time or as closures.

print(model.total.dependencies)
print(model.upstream("total"))
print(model.downstream("theta"))
print(model.graph().to_dot())

('scaled',)
('mu', 'tau', 'theta', 'scaled')
('scaled', 'total', 'spread')
digraph {
  "mu";
  "tau";
  "theta";
  "scaled";
  "total";
  "spread";
  "mu" -> "scaled";
  "tau" -> "scaled";
  "theta" -> "scaled";
  "scaled" -> "total";
  "theta" -> "spread";
}

A slice keeps only what total needs.

print(model.slice("total").source)

mu = elementof(reals)
tau = elementof(posreals)
theta = elementof(cartpow(reals, 4))
scaled = mu .+ tau .* theta
total = sum(scaled)

The model as mathematics, in TeX.

for name, tex in model.math().items():
    print(f"{name}: {tex}")

mu: μ \in \overline{\mathbb{R}}
tau: τ \in \left(0, \infty\right]
theta: θ \in {\overline{\mathbb{R}}}^{4}
scaled: \mathrm{scaled}_{i} = μ + τ  θ_{i}
total: \mathrm{total} = \sum_{i} \mathrm{scaled}_{i}
spread: \mathrm{spread} = \sum_{i} θ_{i}  θ_{i}

A query and the operations its StableHLO program uses.

function = model.query("total").compile()
print(function)
print(dict(function.stats()))

<CompiledFunction (point: {mu: float32[], tau: float32[], theta: float32[4]}) -> float32[]>
{'add': 1, 'broadcast_in_dim': 2, 'constant': 1, 'dot_general': 1, 'reduce': 1, 'reshape': 2, 'total': 8}