def test_applyo():
x = var()
assert run(0, x, applyo("add", (1, 2, 3), x)) == (("add", 1, 2, 3), )
assert run(0, x, applyo(x, (1, 2, 3), ("add", 1, 2, 3))) == ("add", )
assert run(0, x, applyo("add", x, ("add", 1, 2, 3))) == ((1, 2, 3), )
a_lv, b_lv, c_lv = var(), var(), var()
from operator import add
assert run(0, c_lv, applyo(add, (1, 2), c_lv)) == (3, )
assert run(0, c_lv, applyo(add, etuple(1, 2), c_lv)) == (3, )
assert run(0, c_lv, applyo(add, a_lv, c_lv)) == (cons(add, a_lv), )
for obj in (
(1, 2, 3),
(add, 1, 2),
[1, 2, 3],
[add, 1, 2],
etuple(1, 2, 3),
etuple(add, 1, 2),
):
o_rator, o_rands = operator(obj), arguments(obj)
assert run(0, a_lv, applyo(o_rator, o_rands,
a_lv)) == (term(o_rator, o_rands), )
# Just acts like `conso` here
assert run(0, a_lv, applyo(o_rator, a_lv, obj)) == (arguments(obj), )
assert run(0, a_lv, applyo(a_lv, o_rands, obj)) == (operator(obj), )
# Just acts like `conso` here, too
assert run(0, c_lv, applyo(a_lv, b_lv, c_lv)) == (cons(a_lv, b_lv), )
# with pytest.raises(ConsError):
assert run(0, a_lv, applyo(a_lv, b_lv, object())) == ()
assert run(0, a_lv, applyo(1, 2, a_lv)) == ()
def test_assoccomm():
x, a, b, c = tt.dvectors("xabc")
test_expr = x + 1
q = var()
res = run(1, q, applyo(tt.add, etuple(*test_expr.owner.inputs), test_expr))
assert q == res[0]
res = run(1, q, applyo(q, etuple(*test_expr.owner.inputs), test_expr))
assert tt.add == res[0].reify()
res = run(1, q, applyo(tt.add, q, test_expr))
assert mt(tuple(test_expr.owner.inputs)) == res[0]
x = var()
res = run(0, x, eq_comm(mt.mul(a, b), mt.mul(b, x)))
assert (mt(a), ) == res
res = run(0, x, eq_comm(mt.add(a, b), mt.add(b, x)))
assert (mt(a), ) == res
(res, ) = run(0, x, eq_assoc(mt.add(a, b, c), mt.add(a, x)))
assert res == mt(b + c)
(res, ) = run(0, x, eq_assoc(mt.mul(a, b, c), mt.mul(a, x)))
assert res == mt(b * c)
def test_sexp_unify_reify():
"""Make sure we can unify and reify etuples/S-exps."""
# Unify `A . (x + y)`, for `x`, `y` logic variables
A = tf.compat.v1.placeholder(tf.float64,
name="A",
shape=tf.TensorShape([None, None]))
x = tf.compat.v1.placeholder(tf.float64,
name="x",
shape=tf.TensorShape([None, 1]))
y = tf.compat.v1.placeholder(tf.float64,
name="y",
shape=tf.TensorShape([None, 1]))
z = tf.matmul(A, tf.add(x, y))
z_sexp = etuplize(z, shallow=False)
# Let's just be sure that the original TF objects are preserved
assert z_sexp[1].reify() == A
assert z_sexp[2][1].reify() == x
assert z_sexp[2][2].reify() == y
A_lv, x_lv, y_lv = var(), var(), var()
dis_pat = etuple(
TFlowMetaOperator(mt.matmul.op_def, var()),
A_lv,
etuple(TFlowMetaOperator(mt.add.op_def, var()), x_lv, y_lv),
)
s = unify(dis_pat, z_sexp, {})
assert s[A_lv] == mt(A)
assert s[x_lv] == mt(x)
assert s[y_lv] == mt(y)
# Now, we construct a graph that reflects the distributive property and
# reify with the substitutions from the un-distributed form
out_pat = etuple(mt.add, etuple(mt.matmul, A_lv, x_lv),
etuple(mt.matmul, A_lv, y_lv))
z_dist = reify(out_pat, s)
# Evaluate the tuple-expression and get a meta object/graph
z_dist_mt = z_dist.eval_obj
# If all the logic variables were reified, we should be able to
# further reify the meta graph and get a concrete TF graph
z_dist_tf = z_dist_mt.reify()
assert isinstance(z_dist_tf, tf.Tensor)
# Check the first part of `A . x + A . y` (i.e. `A . x`)
assert z_dist_tf.op.inputs[0].op.inputs[0] == A
assert z_dist_tf.op.inputs[0].op.inputs[1] == x
# Now, the second, `A . y`
assert z_dist_tf.op.inputs[1].op.inputs[0] == A
assert z_dist_tf.op.inputs[1].op.inputs[1] == y
def distributes(in_lv, out_lv):
return lall(
# lhs == A * (x + b)
eq(etuple(mt.dot, var("A"), etuple(mt.add, var("x"), var("b"))),
etuplize(in_lv)),
# rhs == A * x + A * b
eq(
etuple(mt.add, etuple(mt.dot, var("A"), var("x")),
etuple(mt.dot, var("A"), var("b"))),
out_lv,
),
)
def distributes(in_lv, out_lv):
A_lv, x_lv, y_lv, z_lv = vars(4)
return lall(
# lhs == A * (x + y + z)
eq_assoccomm(
etuple(_dot, A_lv,
etuple(at.add, x_lv, etuple(at.add, y_lv, z_lv))),
in_lv,
),
# This relation does nothing but provide us with a means of
# generating associative-commutative matches in the `kanren`
# output.
eq((A_lv, x_lv, y_lv, z_lv), out_lv),
)
def test_ConstrainedVar():
cvar = ConstrainedVar(lambda x: isinstance(x, str))
assert repr(cvar).startswith("ConstrainedVar(")
assert repr(cvar).endswith(f", {cvar.token})")
s = unify(cvar, 1)
assert s is False
s = unify(1, cvar)
assert s is False
s = unify(cvar, "hi")
assert s[cvar] == "hi"
s = unify("hi", cvar)
assert s[cvar] == "hi"
x_lv = var()
s = unify(cvar, x_lv)
assert s == {cvar: x_lv}
s = unify(cvar, x_lv, {x_lv: "hi"})
assert s[cvar] == "hi"
s = unify(x_lv, cvar, {x_lv: "hi"})
assert s[cvar] == "hi"
s_orig = {cvar: "hi", x_lv: "hi"}
s = unify(x_lv, cvar, s_orig)
assert s == s_orig
s_orig = {cvar: "hi", x_lv: "bye"}
s = unify(x_lv, cvar, s_orig)
assert s is False
x_at = at.vector("x")
y_at = at.vector("y")
op1_np = CustomOpNoProps(1)
r_at = etuple(op1_np, x_at, y_at)
def constraint(x):
return isinstance(x, tuple)
a_lv = ConstrainedVar(constraint)
res = reify(etuple(op1_np, a_lv), {a_lv: r_at})
assert res[1] == r_at
def cdr_Variable(x):
if x.owner:
x_e = etuple(_car(x), *x.owner.inputs, evaled_obj=x)
else:
raise ConsError("Not a cons pair.")
return x_e[1:]
def test_etuples():
x_at = at.vector("x")
y_at = at.vector("y")
z_at = etuple(x_at, y_at)
res = apply(at.add, z_at)
assert res.owner.op == at.add
assert res.owner.inputs == [x_at, y_at]
w_at = etuple(at.add, x_at, y_at)
res = w_at.evaled_obj
assert res.owner.op == at.add
assert res.owner.inputs == [x_at, y_at]
# This `Op` doesn't expand into an `etuple` (i.e. it's "atomic")
op1_np = CustomOpNoProps(1)
res = apply(op1_np, z_at)
assert res.owner.op == op1_np
q_at = op1_np(x_at, y_at)
res = etuplize(q_at)
assert res[0] == op1_np
with pytest.raises(TypeError):
etuplize(op1_np)
class MyMultiOutOp(Op):
def make_node(self, *inputs):
outputs = [MyType()(), MyType()()]
return Apply(self, list(inputs), outputs)
def perform(self, node, inputs, outputs):
outputs[0] = np.array(inputs[0])
outputs[1] = np.array(inputs[0])
x_at = at.vector("x")
op1_np = MyMultiOutOp()
res = apply(op1_np, etuple(x_at))
assert len(res) == 2
assert res[0].owner.op == op1_np
assert res[1].owner.op == op1_np
def cdr_Op(x):
if not hasattr(x, "__props__"):
raise ConsError("Not a cons pair.")
x_e = etuple(
_car(x),
*[getattr(x, p) for p in getattr(x, "__props__", ())],
evaled_obj=x,
)
return x_e[1:]
def test_kanren_basic():
A_at = at.matrix("A")
x_at = at.vector("x")
y_at = at.dot(A_at, x_at)
q = var()
res = list(run(None, q, eq(y_at, etuple(_dot, q, x_at))))
assert res == [A_at]
def cdr_MetaVariable(x):
"""Return the arguments/tail/CDR of a variable object.
See `cdr_MetaSymbol`
"""
try:
x_e = etuple(_car(x), *x.base_arguments, eval_obj=x)
except NotImplementedError:
raise ConsError("Not a cons pair.")
return x_e[1:]
def test_convert_strs_to_vars():
res = convert_strs_to_vars("a")
assert isinstance(res, Var)
assert res.token == "a"
x_at = at.vector()
y_at = at.vector()
res = convert_strs_to_vars((("a", x_at), y_at))
assert res == etuple(etuple(var("a"), x_at), y_at)
def constraint(x):
return isinstance(x, str)
res = convert_strs_to_vars((({
"pattern": "a",
"constraint": constraint
}, x_at), y_at))
assert res == etuple(etuple(ConstrainedVar(constraint, "a"), x_at), y_at)
# Make sure constrained logic variables are the same across distinct uses
# of their string names
res = convert_strs_to_vars(({
"pattern": "a",
"constraint": constraint
}, "a"))
assert res[0] is res[1]
var_map = {"a": var("a")}
res = convert_strs_to_vars(("a", ), var_map=var_map)
assert res[0] is var_map["a"]
# Make sure numbers and NumPy arrays are converted
val = np.r_[1, 2]
res = convert_strs_to_vars((val, ))
assert isinstance(res[0], Constant)
assert np.array_equal(res[0].data, val)
def normal_normal_regression(Y, X, beta, Y_args_tail=None, beta_args=None):
"""Create a goal for a normal-normal regression of the form `Y ~ N(X * beta, sd**2)`."""
Y_args_tail = Y_args_tail or var()
beta_args = beta_args or var()
Y_args, Y_mean_lv = var(), var()
res = lall(
# `Y` is a `NormalRV`
applyo(mt.NormalRV, Y_args, Y),
# `beta` is also a `NormalRV`
applyo(mt.NormalRV, beta_args, beta),
# Obtain its mean parameter and remaining args
applyo(Y_mean_lv, Y_args_tail, Y_args),
# Relate it to a dot product of `X` and `beta`
applyo(mt.dot, etuple(X, beta), Y_mean_lv),
)
return res
def test_unify_Type():
t1 = TensorType(np.float64, (True, False))
t2 = TensorType(np.float64, (True, False))
# `Type`, `Type`
s = unify(t1, t2)
assert s == {}
# `Type`, `ExpressionTuple`
s = unify(t1, etuple(TensorType, "float64", (1, None)))
assert s == {}
from aesara.scalar.basic import Scalar
st1 = Scalar(np.float64)
st2 = Scalar(np.float64)
s = unify(st1, st2)
assert s == {}
def _convert(y):
if isinstance(y, str):
v = var_map.get(y, var(y))
var_map[y] = v
return v
elif isinstance(y, dict):
pattern = y["pattern"]
if not isinstance(pattern, str):
raise TypeError(
"Constraints can only be assigned to logic variables (i.e. strings)"
)
constraint = y["constraint"]
v = var_map.get(pattern, ConstrainedVar(constraint, pattern))
var_map[pattern] = v
return v
elif isinstance(y, tuple):
return etuple(*tuple(_convert(e) for e in y))
elif isinstance(y, (Number, np.ndarray)):
from aesara.tensor import as_tensor_variable
return as_tensor_variable(y)
return y
def test_etuple_term():
"""Test `etuplize` and `etuple` interaction with `term`."""
# Take apart an already constructed/evaluated meta
# object.
e2 = mt.add(mt.vector(), mt.vector())
e2_et = etuplize(e2)
assert isinstance(e2_et, ExpressionTuple)
# e2_et_expect = etuple(
# mt.add,
# etuple(mt.TensorVariable,
# etuple(mt.TensorType,
# 'float64', (False,), None),
# None, None, None),
# etuple(mt.TensorVariable,
# etuple(mt.TensorType,
# 'float64', (False,), None),
# None, None, None),
# )
e2_et_expect = etuple(mt.add, e2.base_arguments[0], e2.base_arguments[1])
assert e2_et == e2_et_expect
assert e2_et.eval_obj is e2
# Make sure expression expansion works from Theano objects, too.
# First, do it manually.
tt_expr = tt.vector() + tt.vector()
mt_expr = mt(tt_expr)
assert mt_expr.obj is tt_expr
assert mt_expr.reify() is tt_expr
e3 = etuplize(mt_expr)
assert e3 == e2_et
assert e3.eval_obj is mt_expr
assert e3.eval_obj.reify() is tt_expr
# Now, through `etuplize`
e2_et_2 = etuplize(tt_expr)
assert e2_et_2 == e3 == e2_et
assert isinstance(e2_et_2, ExpressionTuple)
assert e2_et_2.eval_obj == tt_expr
test_expr = mt(tt.vector("z") * 7)
assert rator(test_expr) == mt.mul
assert rands(test_expr)[0] == mt(tt.vector("z"))
dim_shuffle_op = rator(rands(test_expr)[1])
assert isinstance(dim_shuffle_op, mt.DimShuffle)
assert rands(rands(test_expr)[1]) == etuple(mt(7))
with pytest.raises(ConsError):
rator(dim_shuffle_op)
# assert rator(dim_shuffle_op) == mt.DimShuffle
# assert rands(dim_shuffle_op) == etuple((), ("x",), True)
const_tensor = rands(rands(test_expr)[1])[0]
with pytest.raises(ConsError):
rator(const_tensor)
with pytest.raises(ConsError):
rands(const_tensor)
et_expr = etuplize(test_expr)
exp_res = etuple(mt.mul, mt(tt.vector("z")),
etuple(mt.DimShuffle((), ("x", ), True), mt(7))
# etuple(etuple(mt.DimShuffle, (), ("x",), True), mt(7))
)
assert et_expr == exp_res
assert exp_res.eval_obj == test_expr
def cdr_Type(x):
x_e = etuple(_car(x),
*[getattr(x, p) for p in getattr(x, "__props__", ())],
evaled_obj=x)
return x_e[1:]
def test_unify_Variable():
x_at = at.vector("x")
y_at = at.vector("y")
z_at = x_at + y_at
# `Variable`, `Variable`
s = unify(z_at, z_at)
assert s == {}
# These `Variable`s have no owners
v1 = MyType()()
v2 = MyType()()
assert v1 != v2
s = unify(v1, v2)
assert s is False
op_lv = var()
z_pat_et = etuple(op_lv, x_at, y_at)
# `Variable`, `ExpressionTuple`
s = unify(z_at, z_pat_et, {})
assert op_lv in s
assert s[op_lv] == z_at.owner.op
res = reify(z_pat_et, s)
assert isinstance(res, ExpressionTuple)
assert equal_computations([res.evaled_obj], [z_at])
z_et = etuple(at.add, x_at, y_at)
# `ExpressionTuple`, `ExpressionTuple`
s = unify(z_et, z_pat_et, {})
assert op_lv in s
assert s[op_lv] == z_et[0]
res = reify(z_pat_et, s)
assert isinstance(res, ExpressionTuple)
assert equal_computations([res.evaled_obj], [z_et.evaled_obj])
# `ExpressionTuple`, `Variable`
s = unify(z_et, x_at, {})
assert s is False
# This `Op` doesn't expand into an `ExpressionTuple`
op1_np = CustomOpNoProps(1)
q_at = op1_np(x_at, y_at)
a_lv = var()
b_lv = var()
# `Variable`, `ExpressionTuple`
s = unify(q_at, etuple(op1_np, a_lv, b_lv))
assert s[a_lv] == x_at
assert s[b_lv] == y_at
def scale_loc_transform(in_expr, out_expr):
"""Create relations for lifting and sinking scale and location parameters of distributions.
I.e. f(a + b*x) -> a + b * f(x)
For example, `in_expr`: f(a + b*x) == `out_expr`: a + b * f(x).
TODO: Match larger distribution families and perform transforms from there.
XXX: PyMC3 rescaling issue (?) doesn't allow us to take the more general
approach, which involves separate scale and location rewrites.
"""
# Scale and location transform expression "pattern" for a Normal term.
normal_mt = mt.NormalRV(var(), var(), size=var(), rng=var(), name=var())
n_name_lv = normal_mt.name
n_mean_lv, n_sd_lv, n_size_lv, n_rng_lv = normal_mt.owner.inputs
offset_name_mt = var()
rct_norm_offset_mt = etuple(
mt.add,
n_mean_lv,
etuple(
mt.mul,
n_sd_lv,
mt.NormalRV(0.0,
1.0,
size=n_size_lv,
rng=n_rng_lv,
name=offset_name_mt),
),
)
# Scale and location transform expression "pattern" for a Cauchy term.
cauchy_mt = mt.CauchyRV(var(), var(), size=var(), rng=var(), name=var())
c_name_lv = cauchy_mt.name
c_mean_lv, c_beta_lv, c_size_lv, c_rng_lv = cauchy_mt.owner.inputs
rct_cauchy_offset_mt = etuple(
mt.add,
c_mean_lv,
etuple(
mt.mul,
c_beta_lv,
mt.CauchyRV(0.0,
1.0,
size=c_size_lv,
rng=c_rng_lv,
name=offset_name_mt),
),
)
# TODO:
# uniform_mt = mt.UniformRV(var(), var(), size=var(), rng=var(), name=var())
# u_name_lv = uniform_mt.name
# u_a_lv, u_b_lv, u_size_lv, u_rng_lv = uniform_mt.owner.inputs
# rct_uniform_scale_mt = etuple(
# mt.mul,
# u_b_lv,
# mt.UniformRV(0.0, 1.0, size=u_size_lv, rng=u_rng_lv, name=offset_name_mt),
# )
# rct_uniform_loc_mt = etuple(mt.add, u_c_lv,
# mt.UniformRV(u_a_lv, u_b_lv,
# size=u_size_lv,
# rng=u_rng_lv,
# name=offset_name_mt))
rels = conde(
[
eq(in_expr, normal_mt),
constant_neq(n_sd_lv, 1),
constant_neq(n_mean_lv, 0),
eq(out_expr, rct_norm_offset_mt),
concat(n_name_lv, "_offset", offset_name_mt),
],
[
eq(in_expr, cauchy_mt),
constant_neq(c_beta_lv, 1),
# TODO: Add a positivity constraint for the scale.
constant_neq(c_mean_lv, 0),
eq(out_expr, rct_cauchy_offset_mt),
concat(c_name_lv, "_offset", offset_name_mt),
],
# TODO:
# [eq(in_expr, uniform_mt),
# lall(
# constant_eq(u_a_lv, 0),
# eq(out_expr, rct_uniform_scale_mt),
# concat(u_name_lv, "_scale", offset_name_mt),
# )],
)
return rels
def test_normal_normal_regression():
tt.config.compute_test_value = "ignore"
theano.config.cxx = ""
np.random.seed(9283)
N = 10
M = 3
a_tt = tt.vector("a")
R_tt = tt.vector("R")
X_tt = tt.matrix("X")
V_tt = tt.vector("V")
a_tt.tag.test_value = np.random.normal(size=M)
R_tt.tag.test_value = np.abs(np.random.normal(size=M))
X = np.random.normal(10, 1, size=N)
X = np.c_[np.ones(10), X, X * X]
X_tt.tag.test_value = X
V_tt.tag.test_value = np.ones(N)
beta_rv = NormalRV(a_tt, R_tt, name="\\beta")
E_y_rv = X_tt.dot(beta_rv)
E_y_rv.name = "E_y"
Y_rv = NormalRV(E_y_rv, V_tt, name="Y")
y_tt = tt.as_tensor_variable(Y_rv.tag.test_value)
y_tt.name = "y"
y_obs_rv = observed(y_tt, Y_rv)
y_obs_rv.name = "y_obs"
#
# Use the relation with identify/match `Y`, `X` and `beta`.
#
y_args_tail_lv, b_args_tail_lv = var(), var()
beta_lv = var()
y_args_lv, y_lv, Y_lv, X_lv = var(), var(), var(), var()
(res, ) = run(
1,
(beta_lv, y_args_tail_lv, b_args_tail_lv),
applyo(mt.observed, y_args_lv, y_obs_rv),
eq(y_args_lv, (y_lv, Y_lv)),
normal_normal_regression(Y_lv, X_lv, beta_lv, y_args_tail_lv,
b_args_tail_lv),
)
# TODO FIXME: This would work if non-op parameters (e.g. names) were covered by
# `operator`/`car`. See `TheanoMetaOperator`.
assert res[0].eval_obj.obj == beta_rv
assert res[0] == etuplize(beta_rv)
assert res[1] == etuplize(Y_rv)[2:]
assert res[2] == etuplize(beta_rv)[1:]
#
# Use the relation with to produce `Y` from given `X` and `beta`.
#
X_new_mt = mt(tt.eye(N, M))
beta_new_mt = mt(NormalRV(0, 1, size=M))
Y_args_cdr_mt = etuplize(Y_rv)[2:]
Y_lv = var()
(res, ) = run(
1, Y_lv,
normal_normal_regression(Y_lv, X_new_mt, beta_new_mt, Y_args_cdr_mt))
Y_out_mt = res.eval_obj
Y_new_mt = etuple(mt.NormalRV, mt.dot(X_new_mt,
beta_new_mt)) + Y_args_cdr_mt
Y_new_mt = Y_new_mt.eval_obj
assert Y_out_mt == Y_new_mt
def normal_qr_transform(in_expr, out_expr):
"""Produce a relation for normal-normal regression and its QR-reduced form.
TODO XXX: This isn't entirely correct (e.g. it needs to also
transform the variance terms), but it demonstrates all the requisite
functionality for this kind of model reformulation.
"""
y_lv, Y_lv, X_lv, beta_lv = var(), var(), var(), var()
Y_args_lv, beta_args_lv = var(), var()
QR_lv, Q_lv, R_lv = var(), var(), var()
beta_til_lv, beta_new_lv = var(), var()
beta_mean_lv, beta_sd_lv = var(), var()
beta_size_lv, beta_rng_lv = var(), var()
Y_new_lv = var()
X_op_lv = var()
in_expr = etuplize(in_expr)
res = lall(
# Only applies to regression models on observed RVs
eq(in_expr, etuple(mt.observed, y_lv, Y_lv)),
# Relate the model components
normal_normal_regression(Y_lv, X_lv, beta_lv, Y_args_lv, beta_args_lv),
# Let's not do all this to an already QR-reduce graph;
# otherwise, we'll loop forever!
applyo(X_op_lv, var(), X_lv),
# XXX: This type of dis-equality goal isn't the best,
# but it will definitely work for now.
neq(mt.nlinalg.qr_full, X_op_lv),
# Relate terms for the QR decomposition
eq(QR_lv, etuple(mt.nlinalg.qr_full, X_lv)),
eq(Q_lv, etuple(itemgetter(0), QR_lv)),
eq(R_lv, etuple(itemgetter(1), QR_lv)),
# The new `beta_tilde`
eq(beta_args_lv, (beta_mean_lv, beta_sd_lv, beta_size_lv, beta_rng_lv)),
eq(
beta_til_lv,
etuple(
mt.NormalRV,
# Use these `tt.[ones|zeros]_like` functions to preserve the
# correct shape (and a valid `tt.dot`).
etuple(mt.zeros_like, beta_mean_lv),
etuple(mt.ones_like, beta_sd_lv),
beta_size_lv,
beta_rng_lv,
),
),
# Relate the new and old coeffs
eq(beta_new_lv, etuple(mt.dot, etuple(mt.nlinalg.matrix_inverse, R_lv), beta_til_lv)),
# Use the relation the other way to produce the new/transformed
# observation distribution
normal_normal_regression(Y_new_lv, Q_lv, beta_til_lv, Y_args_lv),
eq(
out_expr,
[
(
in_expr,
etuple(mt.observed, y_lv, etuple(update_name_suffix, Y_new_lv, Y_lv, "")),
),
(beta_lv, beta_new_lv),
],
),
)
return res
