BUG: `to_graphviz` breaks when the model has a `pm.CustomDist`

[elc45]

Describe the issue:

Title

Reproduceable code example:

import pymc as pm
import pytensor.tensor as pt

def my_logp(value, mu, sigma):
    return -0.5 * pt.log(2 * pt.pi * sigma**2) - 0.5 * ((value - mu) / sigma) ** 2

with pm.Model() as model:
    mu = pm.Normal("mu", 0, 1)
    sigma = pm.Exponential("sigma", 1)
    y = pm.CustomDist(
        "y",
        mu,
        sigma,
        logp=my_logp,
    )

model.to_graphviz()

Error message:

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/link/utils.py:197, in streamline.<locals>.streamline_default_f()
    194 for thunk, node, old_storage in zip(
    195     thunks, order, post_thunk_old_storage, strict=False
    196 ):
--> 197     thunk()
    198     for old_s in old_storage:

File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/graph/op.py:545, in Op.make_py_thunk.<locals>.rval(p, i, o, n, cm)
    537 @is_thunk_type
    538 def rval(
    539     p=p,
   (...)    543     cm=node_compute_map,
    544 ):
--> 545     r = p(n, [x[0] for x in i], o)
    546     for entry in cm:

File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/tensor/random/op.py:430, in RandomVariable.perform(self, node, inputs, outputs)
    428 outputs[0][0] = rng
    429 outputs[1][0] = np.asarray(
--> 430     self.rng_fn(rng, *args, None if size is None else tuple(size)),
    431     dtype=self.dtype,
    432 )

File ~/PyMC/pymc/pymc/distributions/custom.py:87, in CustomDistRV.rng_fn(cls, rng, *args)
     86 size = args.pop(-1)
---> 87 return cls._random_fn(*args, rng=rng, size=size)

File ~/PyMC/pymc/pymc/distributions/custom.py:53, in default_not_implemented.<locals>.func(*args, **kwargs)
     52 def func(*args, **kwargs):
---> 53     raise NotImplementedError(message)

NotImplementedError: Attempted to run random on the CustomDist 'CustomDist_y', but this method had not been provided when the distribution was constructed. Please re-build your model and provide a callable to 'CustomDist_y's random keyword argument.


During handling of the above exception, another exception occurred:

NotImplementedError                       Traceback (most recent call last)
Cell In[2], line 1
----> 1 model.to_graphviz()

File ~/PyMC/pymc/pymc/model/core.py:2034, in Model.to_graphviz(self, var_names, formatting, save, figsize, dpi)
   1974 """Produce a graphviz Digraph from a PyMC model.
   1975 
   1976 Requires graphviz, which may be installed most easily with
   (...)   2030     schools.to_graphviz().render("schools")
   2031 """
   2032 from pymc.model_graph import model_to_graphviz
-> 2034 return model_to_graphviz(
   2035     model=self,
   2036     var_names=var_names,
   2037     formatting=formatting,
   2038     save=save,
   2039     figsize=figsize,
   2040     dpi=dpi,
   2041 )

File ~/PyMC/pymc/pymc/model_graph.py:763, in model_to_graphviz(model, var_names, formatting, save, figsize, dpi, node_formatters, graph_attr, include_dim_lengths)
    759 model = modelcontext(model)
    760 graph = ModelGraph(model)
    761 return make_graph(
    762     model.name,
--> 763     plates=graph.get_plates(var_names=var_names),
    764     edges=graph.edges(var_names=var_names),
    765     formatting=formatting,
    766     save=save,
    767     figsize=figsize,
    768     dpi=dpi,
    769     graph_attr=graph_attr,
    770     node_formatters=node_formatters,
    771     create_plate_label=create_plate_label_with_dim_length
    772     if include_dim_lengths
    773     else create_plate_label_without_dim_length,
    774 )

File ~/PyMC/pymc/pymc/model_graph.py:349, in ModelGraph.get_plates(self, var_names)
    342 # TODO: Evaluate all RV shapes at once
    343 #       This should help find discrepancies, and
    344 #       avoids unnecessary function compiles for determining labels.
    345 dim_lengths: dict[str, int] = {
    346     dim_name: fast_eval(value).item() for dim_name, value in self.model.dim_lengths.items()
    347 }
    348 var_shapes: dict[str, tuple[int, ...]] = {
--> 349     var_name: tuple(map(int, fast_eval(self.model[var_name].shape)))
    350     for var_name in self.vars_to_plot(var_names)
    351 }
    353 for var_name in self.vars_to_plot(var_names):
    354     shape: tuple[int, ...] = var_shapes[var_name]

File ~/PyMC/pymc/pymc/model_graph.py:77, in fast_eval(var)
     76 def fast_eval(var):
---> 77     return function([], var, mode=_cheap_eval_mode)()

File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/compile/function/types.py:1038, in Function.__call__(self, output_subset, *args, **kwargs)
   1036     t0_fn = time.perf_counter()
   1037 try:
-> 1038     outputs = vm() if output_subset is None else vm(output_subset=output_subset)
   1039 except Exception:
   1040     self._restore_defaults()

File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/link/utils.py:201, in streamline.<locals>.streamline_default_f()
    199             old_s[0] = None
    200 except Exception:
--> 201     raise_with_op(fgraph, node, thunk)

File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/link/utils.py:526, in raise_with_op(fgraph, node, thunk, exc_info, storage_map)
    521     warnings.warn(
    522         f"{exc_type} error does not allow us to add an extra error message"
    523     )
    524     # Some exception need extra parameter in inputs. So forget the
    525     # extra long error message in that case.
--> 526 raise exc_value.with_traceback(exc_trace)

File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/link/utils.py:197, in streamline.<locals>.streamline_default_f()
    192 try:
    193     # strict=False because we are in a hot loop
    194     for thunk, node, old_storage in zip(
    195         thunks, order, post_thunk_old_storage, strict=False
    196     ):
--> 197         thunk()
    198         for old_s in old_storage:
    199             old_s[0] = None

File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/graph/op.py:545, in Op.make_py_thunk.<locals>.rval(p, i, o, n, cm)
    537 @is_thunk_type
    538 def rval(
    539     p=p,
   (...)    543     cm=node_compute_map,
    544 ):
--> 545     r = p(n, [x[0] for x in i], o)
    546     for entry in cm:
    547         entry[0] = True

File ~/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/pytensor/tensor/random/op.py:430, in RandomVariable.perform(self, node, inputs, outputs)
    426     rng = custom_rng_deepcopy(rng)
    428 outputs[0][0] = rng
    429 outputs[1][0] = np.asarray(
--> 430     self.rng_fn(rng, *args, None if size is None else tuple(size)),
    431     dtype=self.dtype,
    432 )

File ~/PyMC/pymc/pymc/distributions/custom.py:87, in CustomDistRV.rng_fn(cls, rng, *args)
     85 args = list(args)
     86 size = args.pop(-1)
---> 87 return cls._random_fn(*args, rng=rng, size=size)

File ~/PyMC/pymc/pymc/distributions/custom.py:53, in default_not_implemented.<locals>.func(*args, **kwargs)
     52 def func(*args, **kwargs):
---> 53     raise NotImplementedError(message)

NotImplementedError: Attempted to run random on the CustomDist 'CustomDist_y', but this method had not been provided when the distribution was constructed. Please re-build your model and provide a callable to 'CustomDist_y's random keyword argument.

Apply node that caused the error: CustomDist_y_rv{"(),()->()"}(RNG(<Generator(PCG64) at 0x129C54AC0>), NoneConst{None}, mu, sigma)
Toposort index: 3
Inputs types: [RandomGeneratorType, <pytensor.tensor.type_other.NoneTypeT object at 0x10f97f0e0>, TensorType(float64, shape=()), TensorType(float64, shape=())]
Inputs shapes: ['No shapes', 'No shapes', (), ()]
Inputs strides: ['No strides', 'No strides', (), ()]
Inputs values: [Generator(PCG64) at 0x129C54AC0, None, array(-1.42712261), array(2.46658558)]
Outputs clients: [[], [Shape(y)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/Users/eliotcarlson/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/IPython/core/interactiveshell.py", line 3641, in run_ast_nodes
    if await self.run_code(code, result, async_=asy):
  File "/Users/eliotcarlson/miniforge3/envs/pymc-dev/lib/python3.14/site-packages/IPython/core/interactiveshell.py", line 3701, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/var/folders/sx/5pzrrrlx22v899kvqb67jtvr0000gn/T/ipykernel_81033/3000196339.py", line 12, in <module>
    y = pm.CustomDist(
  File "/Users/eliotcarlson/PyMC/pymc/pymc/distributions/custom.py", line 743, in __new__
    return _CustomDist(
  File "/Users/eliotcarlson/PyMC/pymc/pymc/distributions/distribution.py", line 540, in __new__
    rv_out = cls.dist(*args, **kwargs)
  File "/Users/eliotcarlson/PyMC/pymc/pymc/distributions/custom.py", line 132, in dist
    return super().dist(
  File "/Users/eliotcarlson/PyMC/pymc/pymc/distributions/distribution.py", line 609, in dist
    return cls.rv_op(*dist_params, size=create_size, **kwargs)
  File "/Users/eliotcarlson/PyMC/pymc/pymc/distributions/custom.py", line 191, in rv_op
    return rv_op(*dist_params, **kwargs)

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

PyMC version information:

pymc 5.25.1
pytensor 2.36.3

Context for the issue:

No response