aeppl icon indicating copy to clipboard operation
aeppl copied to clipboard

Implement basic location-scale rewrites

Open brandonwillard opened this issue 1 year ago • 0 comments

We can simplify our IR by using rewrites to derive the exact distributions of location-scale/affine transformations of basic random variables.

For example, the following produces an unnecessarily complicated MeasurableElemwiseTransform instead of a simple NormalRV with transformed parameters:

import aesara
import aesara.tensor as at

from aeppl.rewriting import construct_ir_fgraph


srng = at.random.RandomStream(0)

X_rv = srng.normal(1.0, name="X")

A = at.matrix("A")
b = at.vector("b")

Z_rv = A @ X_rv + b
Z_rv.name = "Z"

z_vv = Z_rv.clone()

fgraph, _, _ = construct_ir_fgraph({Z_rv: z_vv})

aesara.dprint(fgraph)
# ValuedVariable [id A] 4
#  |MeasurableElemwiseTransform{add} [id B] 'Z' 3
#  | |UnmeasurableMeasurableElemwiseTransform{mul} [id C] 2
#  | | |normal_rv{0, (0, 0), floatX, False}.1 [id D] 'X_lifted' 1
#  | | | |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FBE0D3CCAC0>) [id E]
#  | | | |TensorConstant{[]} [id F]
#  | | | |TensorConstant{11} [id G]
#  | | | |TensorConstant{(1, 1) of 1.0} [id H]
#  | | | |TensorConstant{(1, 1) of 1.0} [id H]
#  | | |A [id I]
#  | |InplaceDimShuffle{x,0} [id J] 0
#  |   |b [id K]
#  |Z [id L]

brandonwillard avatar Apr 19 '23 00:04 brandonwillard