Latex representation broken for models with SymbolicDistributions

[ricardoV94]

with pm.Model() as model:
    x = pm.ZeroInflatedPoisson("x", 0.5, 5)
print(model)

  File "/home/ricardo/miniconda3/envs/roche_preclinical_hd/lib/python3.9/site-packages/IPython/core/formatters.py", line 224, in catch_format_error
    r = method(self, *args, **kwargs)
  File "/home/ricardo/miniconda3/envs/roche_preclinical_hd/lib/python3.9/site-packages/IPython/core/formatters.py", line 702, in __call__
    printer.pretty(obj)
  File "/home/ricardo/miniconda3/envs/roche_preclinical_hd/lib/python3.9/site-packages/IPython/lib/pretty.py", line 377, in pretty
    return self.type_pprinters[cls](obj, self, cycle)
  File "/home/ricardo/miniconda3/envs/roche_preclinical_hd/lib/python3.9/site-packages/pymc/printing.py", line 185, in _default_repr_pretty
    output = obj.str_repr()
  File "/home/ricardo/miniconda3/envs/roche_preclinical_hd/lib/python3.9/site-packages/pymc/printing.py", line 83, in str_for_model
    names = [s[: s.index("~") - 1] for s in rv_reprs]
  File "/home/ricardo/miniconda3/envs/roche_preclinical_hd/lib/python3.9/site-packages/pymc/printing.py", line 83, in <listcomp>
    names = [s[: s.index("~") - 1] for s in rv_reprs]
ValueError: substring not found

[larryshamalama]

I can have a look into this

[larryshamalama]

Seems like there are two things here that can be fixed: the graphviz and LaTeX printing representation of the model with SymbolicDistributions.

In general, the approach seems to first call rv.owner.op. Because each SymbolicDistribution defines rv_ops in their own way, i.e. not instances of RandomVariable, just looking at rv.owner.op probably won't work, to my understanding. For mixtures and zero-inflated distributions, rv_op.owner.op yields is an instance of MarginalMixtureRV and, for censored distributions, the owner op is a result from applying the Clip function.

Should graphviz and LaTeX printing be manually adapted to each SymbolicDistribution? Or is there a more convenient way to do this?

[ricardoV94]

We should create a dispatch function like we do for logp or get_moment.

For RandomVariables it will do what you mentioned based on the op, but for symbolic distributions we can customize.

This also ensures we have a default catch all when new Symbolic distributions are implemented without a specific representation, instead of failing like in the original issue here.

[ricardoV94]

Alternatively we can check if aeppl printing capabilities could be useful here: https://github.com/aesara-devs/aeppl/blob/main/aeppl/printing.py

I am not familiar with either that code or pymc's. Maybe @brandonwillard can weigh in

[brandonwillard]

Aesara's printing uses dispatch, but not via singledispatch (see symbolic-pymc's TensorFlow printer for an example of one that does).

Anyway, Aesara's printer framework is easily the best approach for any type of printing, because it already has a lot of things covered, and any custom printing will need to cover most of the same ground.

[ricardoV94]

@larryshamalama Any progress on this one?

[larryshamalama]

Not quite 😅 I started PR #5634 which also sat on the back burner... I can give these efforts a push. Did you want to look into this issue?

[ricardoV94]

Would be nice to fix this one, as is a current V4 blocker :)