ENH: Measurable `dot` and `matmul`

[jessegrabowski]

Before

def scaled_mv_normal(mu, cov, R):
    return pm.MvNormal.dist(mu=R @ mu, cov= R @ cov @ R.T)

After

def scaled_mv_normal(mu, cov, R):
    return R @ pm.MvNormal.dist(mu=mu, cov=cov)

Context for the issue:

Currently, the automatic logp inference isn't able to handle dot products or matrix multiplication of a multivariate normal with a deterministic vector/matrix. Since there are closed form results for these cases (it's essentially the same as repeated shifting/scaling/convolution of normals), it seems like it should be possible to have a measurable_dot and measurable_matmul, at least in the MvNormal case?

[ricardoV94]

Sounds good, we should only need a rewrite that converts the second form to the first?

[jessegrabowski]

Yes I think so. Is there a rewrite to look at that I could pattern a PR after?

[ricardoV94]

Yes I think so. Is there a rewrite to look at that I could pattern a PR after?

It should be like some of these standalone rewrites:

pymc/pymc/logprob/transforms.py

Lines 540 to 544 in 827918b

	@node_rewriter([reciprocal])
	def measurable_reciprocal_to_power(fgraph, node):
	"""Convert reciprocal of `MeasurableVariable`s to power."""
	[inp] = node.inputs
	return [pt.pow(inp, -1.0)]

They are only called when the variable in question (in this casve MV @ x)) has to be measured. So you only need to check the dot involves a MvNormal and return the new equivalent MvNormal.

[ricardoV94]

I think this is implemented, the constant case that is