aesara-devs / aeppl Goto Github PK

import aesara.tensor as at

from aesara.ifelse import ifelse

from aeppl.joint_logprob import factorized_joint_logprob


srng = at.random.RandomStream(seed=2320)

I_rv = srng.bernoulli(0.5, name="I")
X_rv = srng.normal(0, 1, name="X")
Y_rv = srng.gamma(0.5, 0.5, name="Y")

Z_rv = ifelse(I_rv, X_rv, Y_rv)
Z_rv.name = "Z"

z_vv = Z_rv.clone()
i_vv = I_rv.clone()
logp_parts = factorized_joint_logprob({Z_rv: z_vv, I_rv: i_vv})

AssertionError: sigma > 0
Apply node that caused the error: Assert{msg='sigma > 0'}(Elemwise{Composite{((i0 + (i1 * sqr(i2))) - log(i3))}}[(0, 2)].0, All.0)

return -inf

aesara.assert_op.Assert = lambda name: (lambda res, *cond: aesara.tensor.switch(
    aesara.tensor.all(aesara.tensor.stack([c.all() for c in cond])), 
    res, 
    -np.inf
))

import aesara.tensor as at
from aeppl import joint_logprob

weights_rv = at.random.dirichlet([1, 1], name="weights")
components = at.stack([at.random.normal(name="c1"), at.random.normal(name="c2")], axis=-1)
mix_rv = at.sum(components * weights_rv, axis=-1)
mix_rv.name = "mix"

# weights can also just be a deterministic, only thing that matters is that it adds up to 1
weights_vv = weights_rv.clone() 
mix_vv = mix_rv.clone()

joint_logprob({weights_rv: weights_vv, mix_rv: mix_vv})

import aesara.tensor as at
import aeppl
from aeppl.transforms import TransformValuesOpt, DEFAULT_TRANSFORM

p_rv = at.random.beta(1, 1, name="p")
p_vv = p_rv.clone()

tr = TransformValuesOpt({p_vv: DEFAULT_TRANSFORM})
logp_dict = aeppl.factorized_joint_logprob({p_rv: p_vv}, extra_rewrites=tr)
logp_dict

assert p_vv in logp_dict  # raises

x_rv = at.random.uniform(0, 1, name="x")
y_rv = at.random.uniform(0, 5, name="y")
z_rv = at.random.uniform(0, 10, name="z")

x = x_rv.clone()
y = y_rv.clone()
z = z_rv.clone()

joint_logprob(
    [x_rv, y_rv, z_rv],
    {x_rv: x, y_rv: y, z_rv: z},
)

joint_logprob(
    at.stack([x_rv, y_rv, z_rv]),
    {x_rv: x, y_rv: y, z_rv: z},
)

import aesara
import aesara.tensor as at
import numpy as np

value = at.matrix("value")

# Forward
log_value = at.log(value)
shift = at.sum(log_value, -1, keepdims=True) / value.shape[-1]
value = log_value[..., :-1] - shift

# Backward
value = at.concatenate([value, -at.sum(value, -1, keepdims=True)])
exp_value_max = at.exp(value - at.max(value, -1, keepdims=True))
result = exp_value_max / at.sum(exp_value_max, -1, keepdims=True)

result = aesara.function([value],[result])

result(np.array([[0.09,0.01,0.9],[0.09,0.01,0.9]])) # Not sure about the exact values

TypeError: Bad input argument to aesara function with name "/tmp/ipykernel_8640/3915568923.py:1" at index 0 (0-based).
Backtrace when that variable is created:

File "/home/kc611/anaconda3/envs/pymcv4/lib/python3.7/site-packages/ipykernel/zmqshell.py", line 532, in run_cell
return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
File "/home/kc611/anaconda3/envs/pymcv4/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 2899, in run_cell
raw_cell, store_history, silent, shell_futures)
File "/home/kc611/anaconda3/envs/pymcv4/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 2944, in _run_cell
return runner(coro)
File "/home/kc611/anaconda3/envs/pymcv4/lib/python3.7/site-packages/IPython/core/async_helpers.py", line 68, in _pseudo_sync_runner
coro.send(None)
File "/home/kc611/anaconda3/envs/pymcv4/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3170, in run_cell_async
interactivity=interactivity, compiler=compiler, result=result)
File "/home/kc611/anaconda3/envs/pymcv4/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3361, in run_ast_nodes
if (await self.run_code(code, result, async_=asy)):
File "/home/kc611/anaconda3/envs/pymcv4/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3441, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "/tmp/ipykernel_8640/4089953573.py", line 9, in <module>
value = at.concatenate([value, -at.sum(value, -1, keepdims=True)])
Non-unit value on shape on a broadcastable dimension.

import aesara
import aesara.tensor as at
import numpy as np

value = at.matrix("value")

# Forward
log_value = at.log(value)
shift = at.sum(log_value, -1, keepdims=True) / value.shape[-1]
value = log_value[..., :-1] - shift

# Backward
value = at.concatenate([value, -at.sum(value, -1, keepdims=True)], axis = -1) # Concatenating along axis = -1
exp_value_max = at.exp(value - at.max(value, -1, keepdims=True))
result = exp_value_max / at.sum(exp_value_max, -1, keepdims=True)

result = aesara.function([value],[result])

result(np.array([[0.09,0.01,0.9],[0.09,0.01,0.9]])) # Not sure about the exact values
# array([[0.23974926, 0.22131646, 0.53893428],
#        [0.23974926, 0.22131646, 0.53893428]])

def truncate(rv, lower, upper):
  # Try to dispatch on specific graph/Op
  try:
    truncated = _truncate(rv)
  except NotImplementedError:

  # Try to use [i]cdf if they are implemented for given RV
    try: 
      icdf = _icdf(rv)
      cdf_lower = at.exp(_logcdf(rv, lower))
      cdf_upper = at.exp(_logcdf(rv, upper))
      uniform = at.random.uniform(cdf_lower, cdf_upper, size=rv.size)
      truncated = icdf(uniform)
    except NotImplementedError:
    
      # Default to slow while scan graph
      ...

  # Wrap truncated in a custom OpFromGraph that can be easily parsed for the logprob component?
  ...

import aesara
import aesara.tensor as at

from aeppl import joint_logprob


Y_rv = at.random.normal(0, 1, size=(2,), name="Y")

y_vv = Y_rv.clone()
y_vv.name = "y"

y_vv_2 = 2 * y_vv
y_vv_2.name = "y_2"

logprob_non_atomic = joint_logprob({Y_rv: y_vv_2})

aesara.dprint([logprob_non_atomic, y_vv_2])
# Sum{acc_dtype=float64} [id A] ''
#  |Assert{msg='sigma > 0'} [id B] 'y_2_logprob'
#    |Elemwise{sub,no_inplace} [id C] ''
#    | |Elemwise{sub,no_inplace} [id D] ''
#    | | |Elemwise{mul,no_inplace} [id E] ''
#    | | | |InplaceDimShuffle{x} [id F] ''
#    | | | | |TensorConstant{-0.5} [id G]
#    | | | |Elemwise{pow,no_inplace} [id H] ''
#    | | |   |Elemwise{true_div,no_inplace} [id I] ''
#    | | |   | |Elemwise{sub,no_inplace} [id J] ''
#    | | |   | | |Elemwise{mul,no_inplace} [id K] 'y_2'
#    | | |   | | | |InplaceDimShuffle{x} [id L] ''
#    | | |   | | | | |TensorConstant{2} [id M]
#    | | |   | | | |y [id N]
#    | | |   | | |InplaceDimShuffle{x} [id O] ''
#    | | |   | |   |TensorConstant{0} [id P]
#    | | |   | |InplaceDimShuffle{x} [id Q] ''
#    | | |   |   |TensorConstant{1} [id R]
#    | | |   |InplaceDimShuffle{x} [id S] ''
#    | | |     |TensorConstant{2} [id T]
#    | | |InplaceDimShuffle{x} [id U] ''
#    | |   |Elemwise{log,no_inplace} [id V] ''
#    | |     |TensorConstant{2.5066282746310002} [id W]
#    | |InplaceDimShuffle{x} [id X] ''
#    |   |Elemwise{log,no_inplace} [id Y] ''
#    |     |TensorConstant{1} [id R]
#    |All [id Z] ''
#      |Elemwise{gt,no_inplace} [id BA] ''
#        |TensorConstant{1} [id R]
#        |TensorConstant{0.0} [id BB]
# Elemwise{mul,no_inplace} [id BC] 'y_2'
#  |InplaceDimShuffle{x} [id BD] ''
#  | |TensorConstant{2} [id M]
#  |y [id N]

import aesara.tensor as at
from aeppl import joint_logprob
from aeppl.transforms import TransformValuesOpt, LogOddsTransform

idx_rv = at.random.bernoulli(0.5, name='idx')
mixture1 = at.random.beta(100, 1, name='mixture_comp1')
mixture2 = at.random.beta(1, 100, name='mixture_comp2')
mixture_rv = at.stack(mixture1, mixture2)[idx_rv]  # MixtureRV is artificially created via `at.stack` to trigger `OpFromGraph`
mixture_rv.name = 'mixture'

idx_vv = idx_rv.clone()
mixture_vv = mixture_rv.clone()

transform_opt = TransformValuesOpt({mixture_vv: LogOddsTransform()})
logp = joint_logprob(
    {mixture_rv: mixture_vv, idx_rv: idx_vv},
    extra_rewrites=transform_opt,
)

logp.eval({idx_vv: 0, mixture_vv: 10}) 
# array(-inf)

idx_rv = at.random.bernoulli(0.5, name='idx')
mixture = at.random.beta([100, 1], [1, 100], name='mixture_components')
mixture_rv = mixture[idx_rv]
mixture_rv.name = 'mixture'

idx_vv = idx_rv.clone()
mixture_vv = mixture_rv.clone()

transform_opt = TransformValuesOpt({mixture_vv: LogOddsTransform()})
logp = joint_logprob(
    {mixture_rv: mixture_vv, idx_rv: idx_vv},
    extra_rewrites=transform_opt,
)

logp.eval({idx_vv: 0, mixture_vv: 10})
# array(-6.09256229)

import aesara.tensor as at
import aeppl

x_rv = at.random.normal(name='x')
y_rv = at.random.normal(x_rv, 1, name="y")

y_vv = y_rv.clone()

logprob_dict = aeppl.factorized_joint_logprob({y_rv: y_vv})
logprob_dict

KeyError                                  Traceback (most recent call last)
<ipython-input-6-2e07a6eebca4> in <module>
      7 y_vv = y_rv.clone()
      8 
----> 9 logprob_dict = aeppl.factorized_joint_logprob({y_rv: y_vv})

~/.local/lib/python3.8/site-packages/aeppl/joint_logprob.py in factorized_joint_logprob(rv_values, warn_missing_rvs, extra_rewrites, **kwargs)
    165             continue
    166 
--> 167         q_rv_value_vars = [
    168             replacements[q_rv_var]
    169             for q_rv_var in outputs

~/.local/lib/python3.8/site-packages/aeppl/joint_logprob.py in <listcomp>(.0)
    166 
    167         q_rv_value_vars = [
--> 168             replacements[q_rv_var]
    169             for q_rv_var in outputs
    170             if not getattr(q_rv_var.tag, "ignore_logprob", False)

KeyError: x

import aesara.tensor as at
import aeppl

x_rv = at.random.normal(name="x")
y_rv = at.cos(x_rv)

z_rv = at.random.normal(name='z')  # Just so that the returned output is not `None`

y_vv = y_rv.clone()
z_vv = z_rv.clone()

aeppl.joint_logprob({y_rv: y_vv, z_rv: z_vv})

missing_value_terms = set(original_values.values()) - set(logprob_vars.keys()) 
if missing_value_terms:
    raise NotImplementedError(
        f"The logprob terms for the following value variables could not be derived: {missing_value_terms}"
    )

    x = at.clip(at.random.normal(loc=1), 0.5, 1.5); x.name = "x"
    y = at.random.beta(1, 2, size=None, name="y")

    comps = at.stack(x, y); comps.name = "comps"
    idxs = at.random.bernoulli(0.4, size=None, name="idxs")
    mix = comps[idxs]; mix.name = "mix"

    mix_vv = mix.clone()
    idxs_vv = idxs.clone()
    logp = joint_logprob({idxs: idxs_vv, mix: mix_vv})

import aesara
import aesara.tensor as at

from aeppl.joint_logprob import factorized_joint_logprob


srng = at.random.RandomStream(seed=2320)
N = 10

Y_rv = srng.normal(0, 1, size=N, name="Y")


def step_fn(y_t):
    return y_t


Y_1T_rv, _ = aesara.scan(
    fn=step_fn,
    sequences=[Y_rv],
    strict=True,
)

y_vv = Y_1T_rv.clone()
y_vv.name = "y"

logp_parts = factorized_joint_logprob({Y_1T_rv: y_vv})

# Apply a rewrite like `local_rv_size_lift` to get properly `size`-broadcasted parameters
# in a new variable `new_Y_rv`
mu_bcast, sigma_bcast = new_Y_rv.owner.inputs[3:]

def new_step_fn(mu_t, sigma_t)
    return Y_rv.owner.op(mu_t, sigma_t, name="Y[t]")

new_Y_1T_rv, _ = aesara.scan(
    fn=new_step_fn,
    sequences=[mu_bcast, sigma_bcast],
    strict=True,
)

loc = at.vector("loc")
y_base_rv = at.random.normal(0, 1, name="y_base_rv", size=2)
y_rv = loc + y_base_rv
y_val = y_rv.clone()

logp = joint_logprob(y_rv, {y_rv: y_val})
logp.eval({y_val: [0, 0, 0, 0], loc_val: [0, 0, 0, 0]})

import aesara
import aesara.tensor as at
from aeppl import joint_logprob

x = at.scalar('x')
beta_rv = at.random.normal(0, 1, name='beta')
y_rv = at.random.normal(beta_rv*x, 1, name='y')

beta = beta_rv.type()
y = y_rv.type()

logp = joint_logprob(y_rv, {beta_rv: beta, y_rv: y})
logp_fun = aesara.function([x, beta, y], logp)

---------------------------------------------------------------------------

UnusedInputError                          Traceback (most recent call last)

<ipython-input-1-ea0e0eb814fa> in <module>
     11 
     12 logp = joint_logprob(y_rv, {beta_rv: beta, y_rv: y})
---> 13 logp_fun = aesara.function([x, beta, y], logp)

~/Documents/Projects/aeppl/venv/lib/python3.8/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,

~/Documents/Projects/aeppl/venv/lib/python3.8/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,

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

~/Documents/Projects/aeppl/venv/lib/python3.8/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)
   1573 
   1574         # Check if some input variables are unused
-> 1575         self._check_unused_inputs(inputs, outputs, on_unused_input)
   1576 
   1577         # Make a list of (SymbolicInput|SymblicInputKits, indices,

~/Documents/Projects/aeppl/venv/lib/python3.8/site-packages/aesara/compile/function/types.py in _check_unused_inputs(self, inputs, outputs, on_unused_input)
   1745                     )
   1746                 elif on_unused_input == "raise":
-> 1747                     raise UnusedInputError(msg % (inputs.index(i), i.variable, err_msg))
   1748                 else:
   1749                     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: x.
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'.

logp = joint_logprob(y_rv, {beta_rv: beta, y_rv: y, x:x})

import aesara
import aesara.tensor as at
import aeppl

a = at.random.normal()
b = at.random.normal(loc=a)

a_val = at.dscalar()
b_val = at.dscalar()

aeppl.factorized_joint_logprob({b: b_val}, sample=[a])

aeppl.factorized_joint_logprob({b: b_val}, marginalize=[a])  # closed form solution only
aeppl.factorized_joint_logprob({b: b_val}, marginalize=[a], integration_options={a: "gauss_hermite(21)"}) # some numerical integration

import aesara
import aesara.tensor as at
from aesara.compile.mode import Mode
from aeppl.transforms import IntervalTransform

a, b, x = at.scalars('a', 'b', 'x')
intvl = IntervalTransform(args_fn=lambda *args: (a, b))
graph = intvl.forward(intvl.backward(x))

graph_fn = aesara.function([a, b, x], graph, mode=Mode().excluding('fusion'))
aesara.dprint(graph_fn)

import aesara.tensor as at

from aeppl.joint_logprob import factorized_joint_logprob


srng = at.random.RandomStream(seed=2320)

I_rv = srng.bernoulli(0.5, size=10, name="I")
X_rv = srng.normal(0, 1, name="X")
Y_rv = srng.gamma(0.5, 0.5, size=10, name="Y")

Z_rv = at.switch(I_rv, X_rv, Y_rv)
Z_rv.name = "Z"

z_vv = Z_rv.clone()
i_vv = I_rv.clone()
logp_parts = factorized_joint_logprob({Z_rv: z_vv, I_rv: i_vv})

from aeppl.transforms import Simplex
import aesara.tensor as at

x = at.fmatrix('x')
assert x.dtype == "float32"
forw_x = Simplex().forward(x)
assert forw_x.dtype == "float32"  # Raises

x_rv = at.random.normal()
potential = at.switch(x_rv < 0, -np.inf, 0)

x = x_rv.clone()

joint_logprob([x_rv, potential], {x_rv: x, potential: None})

import aesara.tensor as at

from aeppl import joint_logprob


a, b = at.scalars("ab")
Y_rv = a * at.random.normal(0, 1, name="Y") + b

y_val = Y_rv.type()
logp = joint_logprob(Y_rv, {Y_rv: y_val})

import aesara
import aesara.tensor as at
from aesara.graph.opt import in2out, out2in

from aeppl import joint_logprob
from aeppl.transforms import transform_logprob

# Does not matter if I set ignore_newtrees to True, or if I use out2in
transform_opt = in2out(transform_logprob, ignore_newtrees=False)

l_rv = at.random.uniform(0, 1, name='l_rv')
u_rv = at.random.uniform(2, 4, name='u_rv')
# l_rv = at.random.normal(name='l_rv')
# u_rv = at.random.normal(name='u_rv')
y_rv = at.random.uniform(l_rv, u_rv, name='y_rv')

l = l_rv.type()
l.name = 'l'
u = u_rv.type()
u.name = 'u'
y = y_rv.type()
y.name='y'

logp = joint_logprob(y_rv, {y_rv:y, l_rv:l, u_rv:u}, extra_rewrites=transform_opt)
logp.eval({l: -1, u:1, y:0}) 

AttributeError: 'numpy.random._generator.Generator' object has no attribute 'uniform_interval'

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'.

import aeppl
import aesara.tensor as at
import aesara
with aesara.config.change_flags(floatX="float32", warn_float64="raise"):
    aeppl.logprob(at.random.normal(), at.constant(1.))

/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aeppl/logprob.py(75)normal_logprob()
-> - at.log(np.sqrt(2.0 * np.pi))
(Pdb) l
 70     def normal_logprob(op, values, *inputs, **kwargs):
 71         (value,) = values
 72         mu, sigma = inputs[3:]
 73         res = (
 74             -0.5 * at.pow((value - mu) / sigma, 2)
 75  ->         - at.log(np.sqrt(2.0 * np.pi))
 76             - at.log(sigma)
 77         )
 78         res = Assert("sigma > 0")(res, at.all(at.gt(sigma, 0.0)))
 79         return res

import aesara.tensor as at
from aeppl import joint_logprob

shift_rv = at.random.normal(0, 1)
y_rv = at.random.exponential(1) + shift

shift = shift_rv.clone()
y = y_rv.clone()

logp = joint_logprob(y_rv, {y_rv: y, shift_rv: shift}, transfoms={y: Interval(shift_rv, None)})

import aesara.tensor.random as atr
from aeppl import joint_logprob
#%%

x_rv = atr.normal(name="x")
y_rv = atr.normal(x_rv, 1, size=2, name="y")

x = x_rv.clone()
y = y_rv.clone()

logprob = joint_logprob([y_rv], {x_rv: x, y_rv: y})

eval1 = logprob.eval({x: 0, y: [0, 0]}).sum()
eval1  # -3.6757541328186907
#%%

x_rv = atr.normal(name="x")
y_rv1 = atr.normal(x_rv, 1, name="y1")
y_rv2 = atr.normal(x_rv, 1, name="y2")

x = x_rv.clone()
y1 = y_rv1.clone()
y2 = y_rv2.clone()

logprob = joint_logprob([y_rv1, y_rv2], {x_rv: x, y_rv1: y1, y_rv2: y2})

eval2 = logprob.eval({x: 0, y1: 0, y2: 0})
eval2  # array(-2.7568156)

x_rv = at.random.bernoulli(0.7)
y_rv = at.normal([0, 1], [1, 1])[x_rv]
y = y_rv.type()

logp = joint_logprob(y_rv, {y_rv: y}, marginalize={x_rv})

y_rv = at.max(at.random.uniform(0, 1, size=3))   # or min
y_rv = at.sort(at.random.uniform(0, 1, size=3))
y_rv = at.sort(at.random.uniform(0, 1, size=3))[idx]  # max / min correspond to idx==-1 or idx==0

y_rv = at.max(at.random.uniform([0, 1], [2, 3]))
y_rv = at.max(at.stack([at.random.uniform(0, 1), at.random.normal(0, 1)]))

import numpy as np

import aesara
import aesara.tensor as at

import aeppl
from aeppl.transforms import TransformValuesOpt, LogOddsTransform

init = at.random.beta(1, 1)
innov, _ = aesara.scan(
    fn=lambda prev_innov: at.random.beta(prev_innov * 10, (1 - prev_innov) * 10),
    outputs_info=[init],
    n_steps=10,
)

init_vv = init.clone()
innov_vv = innov.clone()

tr = TransformValuesOpt({
    init_vv: LogOddsTransform(),
    innov_vv: LogOddsTransform(),
})
logp = aeppl.joint_logprob({init: init_vv, innov: innov_vv}, extra_rewrites=tr)

# init_vv is properly transformed
logp.eval({init_vv: -0.5, innov_vv: np.full(10, 0.5)})
# array(7.22123237)

# innov_vv is not, eval raises ParameterValueError
logp.eval({init_vv: -0.5, innov_vv: np.full(10, -0.5)})

---------------------------------------------------------------------------
ParameterValueError                       Traceback (most recent call last)
scan_perform.pyx in aesara.scan.scan_perform.perform()

ParameterValueError: 0 <= value <= 1, alpha > 0, beta > 0