Broadcasted variable not replaced in downstream terms

[ricardoV94]

Please provide a minimal, self-contained, and reproducible example.

import aesara.tensor as at
import aeppl

cov = at.eye(2)
mu_rv = at.random.normal(size=2, name='mu')
mu_broadcasted = at.broadcast_arrays(mu_rv, cov[..., -1])[0]
# mu_broadcasted = mu_rv  # No error if this line is uncommented
x_rv = at.random.multivariate_normal(mu_broadcasted, cov, name='x')

mu_vv = mu_rv.clone()
x_vv = x_rv.clone()
logp = aeppl.factorized_joint_logprob({mu_rv: mu_vv, x_rv: x_vv})

aeppl/joint_logprob.py:161: UserWarning: Found a random variable that was neither among the observations nor the conditioned variables: normal_rv{0, (0, 0), floatX, False}(RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F51B7F3EC80>), TensorConstant{[]}, TensorConstant{11}, TensorConstant{(2,) of 0.0}, TensorConstant{(2,) of 1.0})
  warnings.warn(

logp[x_vv].eval({mu_vv: [0, 0], x_vv: [0, 0]})

UnusedInputError                          Traceback (most recent call last)
/tmp/ipykernel_55923/2490478877.py in <module>
----> 1 logp[x_vv].eval({mu_vv: [0, 0], x_vv: [0, 0]})

~/miniconda3/envs/aeppl/lib/python3.9/site-packages/aesara/graph/basic.py in eval(self, inputs_to_values)
    548         inputs = tuple(sorted(inputs_to_values.keys(), key=id))
    549         if inputs not in self._fn_cache:
--> 550             self._fn_cache[inputs] = function(inputs, self)
    551         args = [inputs_to_values[param] for param in inputs]
    552 

~/miniconda3/envs/aeppl/lib/python3.9/site-packages/aesara/compile/function/__init__.py in function(inputs, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input)
    335         # note: pfunc will also call orig_function -- orig_function is
    336         #      a choke point that all compilation must pass through
--> 337         fn = pfunc(
    338             params=inputs,
    339             outputs=outputs,

~/miniconda3/envs/aeppl/lib/python3.9/site-packages/aesara/compile/function/pfunc.py in pfunc(params, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input, output_keys)
    522         inputs.append(si)
    523 
--> 524     return orig_function(
    525         inputs,
    526         cloned_outputs,

~/miniconda3/envs/aeppl/lib/python3.9/site-packages/aesara/compile/function/types.py in orig_function(inputs, outputs, mode, accept_inplace, name, profile, on_unused_input, output_keys)
   1971     try:
   1972         Maker = getattr(mode, "function_maker", FunctionMaker)
-> 1973         m = Maker(
   1974             inputs,
   1975             outputs,

~/miniconda3/envs/aeppl/lib/python3.9/site-packages/aesara/compile/function/types.py in __init__(self, inputs, outputs, mode, accept_inplace, function_builder, profile, on_unused_input, fgraph, output_keys, name)
   1574 
   1575         # Check if some input variables are unused
-> 1576         self._check_unused_inputs(inputs, outputs, on_unused_input)
   1577 
   1578         # Make a list of (SymbolicInput|SymblicInputKits, indices,

~/miniconda3/envs/aeppl/lib/python3.9/site-packages/aesara/compile/function/types.py in _check_unused_inputs(self, inputs, outputs, on_unused_input)
   1746                     )
   1747                 elif on_unused_input == "raise":
-> 1748                     raise UnusedInputError(msg % (inputs.index(i), i.variable, err_msg))
   1749                 else:
   1750                     raise ValueError(

UnusedInputError: aesara.function was asked to create a function computing outputs given certain inputs, but the provided input variable at index 0 is not part of the computational graph needed to compute the outputs: mu.
To make this error into a warning, you can pass the parameter on_unused_input='warn' to aesara.function. To disable it completely, use on_unused_input='ignore'.

First seen in pymc-devs/pymc#5382 (comment)

[ricardoV94]

Seems to be due to the naive_bcast_rv_lift, which breaks up the relationship between the original mu_rv and the mu_broadcasted