aeppl
aeppl copied to clipboard
Implement basic location-scale rewrites
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]