Similar to the Dimshuffle case, the best is to let canonicalization move SpecifyShape's out of the way (and improve Aesara if there are obvious missing cases), and only apply a logprob rewrite as the last resort.
The logprob rewrite should be pretty simple
def logprob_specify_shape(op, values, inner_rv, *shapes, **kwargs):
(value,) = values
# transfer specify_shape from rv to value
value = at.specify_shape(value, shapes)
return logprob(inner_rv, value)
Similar to the
Dimshufflecase, the best is to let canonicalization move SpecifyShape's out of the way (and improve Aesara if there are obvious missing cases), and only apply a logprob rewrite as the last resort.The logprob rewrite should be pretty simple