You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

627 lines
20 KiB

"""
Important note on tests in this module - the Aesara printing functions use a
global cache by default, which means that tests using it will modify global
state and thus not be independent from each other. Instead of using the "cache"
keyword argument each time, this module uses the aesara_code_ and
aesara_function_ functions defined below which default to using a new, empty
cache instead.
"""
import logging
from sympy.external import import_module
from sympy.testing.pytest import raises, SKIP
from sympy.utilities.exceptions import ignore_warnings
aesaralogger = logging.getLogger('aesara.configdefaults')
aesaralogger.setLevel(logging.CRITICAL)
aesara = import_module('aesara')
aesaralogger.setLevel(logging.WARNING)
if aesara:
import numpy as np
aet = aesara.tensor
from aesara.scalar.basic import Scalar
from aesara.graph.basic import Variable
from aesara.tensor.var import TensorVariable
from aesara.tensor.elemwise import Elemwise, DimShuffle
from aesara.tensor.math import Dot
from sympy.printing.aesaracode import true_divide
xt, yt, zt = [aet.scalar(name, 'floatX') for name in 'xyz']
Xt, Yt, Zt = [aet.tensor('floatX', (False, False), name=n) for n in 'XYZ']
else:
#bin/test will not execute any tests now
disabled = True
import sympy as sy
from sympy.core.singleton import S
from sympy.abc import x, y, z, t
from sympy.printing.aesaracode import (aesara_code, dim_handling,
aesara_function)
# Default set of matrix symbols for testing - make square so we can both
# multiply and perform elementwise operations between them.
X, Y, Z = [sy.MatrixSymbol(n, 4, 4) for n in 'XYZ']
# For testing AppliedUndef
f_t = sy.Function('f')(t)
def aesara_code_(expr, **kwargs):
""" Wrapper for aesara_code that uses a new, empty cache by default. """
kwargs.setdefault('cache', {})
return aesara_code(expr, **kwargs)
def aesara_function_(inputs, outputs, **kwargs):
""" Wrapper for aesara_function that uses a new, empty cache by default. """
kwargs.setdefault('cache', {})
return aesara_function(inputs, outputs, **kwargs)
def fgraph_of(*exprs):
""" Transform SymPy expressions into Aesara Computation.
Parameters
==========
exprs
SymPy expressions
Returns
=======
aesara.graph.fg.FunctionGraph
"""
outs = list(map(aesara_code_, exprs))
ins = list(aesara.graph.basic.graph_inputs(outs))
ins, outs = aesara.graph.basic.clone(ins, outs)
return aesara.graph.fg.FunctionGraph(ins, outs)
def aesara_simplify(fgraph):
""" Simplify a Aesara Computation.
Parameters
==========
fgraph : aesara.graph.fg.FunctionGraph
Returns
=======
aesara.graph.fg.FunctionGraph
"""
mode = aesara.compile.get_default_mode().excluding("fusion")
fgraph = fgraph.clone()
mode.optimizer.optimize(fgraph)
return fgraph
def theq(a, b):
""" Test two Aesara objects for equality.
Also accepts numeric types and lists/tuples of supported types.
Note - debugprint() has a bug where it will accept numeric types but does
not respect the "file" argument and in this case and instead prints the number
to stdout and returns an empty string. This can lead to tests passing where
they should fail because any two numbers will always compare as equal. To
prevent this we treat numbers as a separate case.
"""
numeric_types = (int, float, np.number)
a_is_num = isinstance(a, numeric_types)
b_is_num = isinstance(b, numeric_types)
# Compare numeric types using regular equality
if a_is_num or b_is_num:
if not (a_is_num and b_is_num):
return False
return a == b
# Compare sequences element-wise
a_is_seq = isinstance(a, (tuple, list))
b_is_seq = isinstance(b, (tuple, list))
if a_is_seq or b_is_seq:
if not (a_is_seq and b_is_seq) or type(a) != type(b):
return False
return list(map(theq, a)) == list(map(theq, b))
# Otherwise, assume debugprint() can handle it
astr = aesara.printing.debugprint(a, file='str')
bstr = aesara.printing.debugprint(b, file='str')
# Check for bug mentioned above
for argname, argval, argstr in [('a', a, astr), ('b', b, bstr)]:
if argstr == '':
raise TypeError(
'aesara.printing.debugprint(%s) returned empty string '
'(%s is instance of %r)'
% (argname, argname, type(argval))
)
return astr == bstr
def test_example_symbols():
"""
Check that the example symbols in this module print to their Aesara
equivalents, as many of the other tests depend on this.
"""
assert theq(xt, aesara_code_(x))
assert theq(yt, aesara_code_(y))
assert theq(zt, aesara_code_(z))
assert theq(Xt, aesara_code_(X))
assert theq(Yt, aesara_code_(Y))
assert theq(Zt, aesara_code_(Z))
def test_Symbol():
""" Test printing a Symbol to a aesara variable. """
xx = aesara_code_(x)
assert isinstance(xx, Variable)
assert xx.broadcastable == ()
assert xx.name == x.name
xx2 = aesara_code_(x, broadcastables={x: (False,)})
assert xx2.broadcastable == (False,)
assert xx2.name == x.name
def test_MatrixSymbol():
""" Test printing a MatrixSymbol to a aesara variable. """
XX = aesara_code_(X)
assert isinstance(XX, TensorVariable)
assert XX.broadcastable == (False, False)
@SKIP # TODO - this is currently not checked but should be implemented
def test_MatrixSymbol_wrong_dims():
""" Test MatrixSymbol with invalid broadcastable. """
bcs = [(), (False,), (True,), (True, False), (False, True,), (True, True)]
for bc in bcs:
with raises(ValueError):
aesara_code_(X, broadcastables={X: bc})
def test_AppliedUndef():
""" Test printing AppliedUndef instance, which works similarly to Symbol. """
ftt = aesara_code_(f_t)
assert isinstance(ftt, TensorVariable)
assert ftt.broadcastable == ()
assert ftt.name == 'f_t'
def test_add():
expr = x + y
comp = aesara_code_(expr)
assert comp.owner.op == aesara.tensor.add
def test_trig():
assert theq(aesara_code_(sy.sin(x)), aet.sin(xt))
assert theq(aesara_code_(sy.tan(x)), aet.tan(xt))
def test_many():
""" Test printing a complex expression with multiple symbols. """
expr = sy.exp(x**2 + sy.cos(y)) * sy.log(2*z)
comp = aesara_code_(expr)
expected = aet.exp(xt**2 + aet.cos(yt)) * aet.log(2*zt)
assert theq(comp, expected)
def test_dtype():
""" Test specifying specific data types through the dtype argument. """
for dtype in ['float32', 'float64', 'int8', 'int16', 'int32', 'int64']:
assert aesara_code_(x, dtypes={x: dtype}).type.dtype == dtype
# "floatX" type
assert aesara_code_(x, dtypes={x: 'floatX'}).type.dtype in ('float32', 'float64')
# Type promotion
assert aesara_code_(x + 1, dtypes={x: 'float32'}).type.dtype == 'float32'
assert aesara_code_(x + y, dtypes={x: 'float64', y: 'float32'}).type.dtype == 'float64'
def test_broadcastables():
""" Test the "broadcastables" argument when printing symbol-like objects. """
# No restrictions on shape
for s in [x, f_t]:
for bc in [(), (False,), (True,), (False, False), (True, False)]:
assert aesara_code_(s, broadcastables={s: bc}).broadcastable == bc
# TODO - matrix broadcasting?
def test_broadcasting():
""" Test "broadcastable" attribute after applying element-wise binary op. """
expr = x + y
cases = [
[(), (), ()],
[(False,), (False,), (False,)],
[(True,), (False,), (False,)],
[(False, True), (False, False), (False, False)],
[(True, False), (False, False), (False, False)],
]
for bc1, bc2, bc3 in cases:
comp = aesara_code_(expr, broadcastables={x: bc1, y: bc2})
assert comp.broadcastable == bc3
def test_MatMul():
expr = X*Y*Z
expr_t = aesara_code_(expr)
assert isinstance(expr_t.owner.op, Dot)
assert theq(expr_t, Xt.dot(Yt).dot(Zt))
def test_Transpose():
assert isinstance(aesara_code_(X.T).owner.op, DimShuffle)
def test_MatAdd():
expr = X+Y+Z
assert isinstance(aesara_code_(expr).owner.op, Elemwise)
def test_Rationals():
assert theq(aesara_code_(sy.Integer(2) / 3), true_divide(2, 3))
assert theq(aesara_code_(S.Half), true_divide(1, 2))
def test_Integers():
assert aesara_code_(sy.Integer(3)) == 3
def test_factorial():
n = sy.Symbol('n')
assert aesara_code_(sy.factorial(n))
def test_Derivative():
with ignore_warnings(UserWarning):
simp = lambda expr: aesara_simplify(fgraph_of(expr))
assert theq(simp(aesara_code_(sy.Derivative(sy.sin(x), x, evaluate=False))),
simp(aesara.grad(aet.sin(xt), xt)))
def test_aesara_function_simple():
""" Test aesara_function() with single output. """
f = aesara_function_([x, y], [x+y])
assert f(2, 3) == 5
def test_aesara_function_multi():
""" Test aesara_function() with multiple outputs. """
f = aesara_function_([x, y], [x+y, x-y])
o1, o2 = f(2, 3)
assert o1 == 5
assert o2 == -1
def test_aesara_function_numpy():
""" Test aesara_function() vs Numpy implementation. """
f = aesara_function_([x, y], [x+y], dim=1,
dtypes={x: 'float64', y: 'float64'})
assert np.linalg.norm(f([1, 2], [3, 4]) - np.asarray([4, 6])) < 1e-9
f = aesara_function_([x, y], [x+y], dtypes={x: 'float64', y: 'float64'},
dim=1)
xx = np.arange(3).astype('float64')
yy = 2*np.arange(3).astype('float64')
assert np.linalg.norm(f(xx, yy) - 3*np.arange(3)) < 1e-9
def test_aesara_function_matrix():
m = sy.Matrix([[x, y], [z, x + y + z]])
expected = np.array([[1.0, 2.0], [3.0, 1.0 + 2.0 + 3.0]])
f = aesara_function_([x, y, z], [m])
np.testing.assert_allclose(f(1.0, 2.0, 3.0), expected)
f = aesara_function_([x, y, z], [m], scalar=True)
np.testing.assert_allclose(f(1.0, 2.0, 3.0), expected)
f = aesara_function_([x, y, z], [m, m])
assert isinstance(f(1.0, 2.0, 3.0), type([]))
np.testing.assert_allclose(f(1.0, 2.0, 3.0)[0], expected)
np.testing.assert_allclose(f(1.0, 2.0, 3.0)[1], expected)
def test_dim_handling():
assert dim_handling([x], dim=2) == {x: (False, False)}
assert dim_handling([x, y], dims={x: 1, y: 2}) == {x: (False, True),
y: (False, False)}
assert dim_handling([x], broadcastables={x: (False,)}) == {x: (False,)}
def test_aesara_function_kwargs():
"""
Test passing additional kwargs from aesara_function() to aesara.function().
"""
import numpy as np
f = aesara_function_([x, y, z], [x+y], dim=1, on_unused_input='ignore',
dtypes={x: 'float64', y: 'float64', z: 'float64'})
assert np.linalg.norm(f([1, 2], [3, 4], [0, 0]) - np.asarray([4, 6])) < 1e-9
f = aesara_function_([x, y, z], [x+y],
dtypes={x: 'float64', y: 'float64', z: 'float64'},
dim=1, on_unused_input='ignore')
xx = np.arange(3).astype('float64')
yy = 2*np.arange(3).astype('float64')
zz = 2*np.arange(3).astype('float64')
assert np.linalg.norm(f(xx, yy, zz) - 3*np.arange(3)) < 1e-9
def test_aesara_function_scalar():
""" Test the "scalar" argument to aesara_function(). """
from aesara.compile.function.types import Function
args = [
([x, y], [x + y], None, [0]), # Single 0d output
([X, Y], [X + Y], None, [2]), # Single 2d output
([x, y], [x + y], {x: 0, y: 1}, [1]), # Single 1d output
([x, y], [x + y, x - y], None, [0, 0]), # Two 0d outputs
([x, y, X, Y], [x + y, X + Y], None, [0, 2]), # One 0d output, one 2d
]
# Create and test functions with and without the scalar setting
for inputs, outputs, in_dims, out_dims in args:
for scalar in [False, True]:
f = aesara_function_(inputs, outputs, dims=in_dims, scalar=scalar)
# Check the aesara_function attribute is set whether wrapped or not
assert isinstance(f.aesara_function, Function)
# Feed in inputs of the appropriate size and get outputs
in_values = [
np.ones([1 if bc else 5 for bc in i.type.broadcastable])
for i in f.aesara_function.input_storage
]
out_values = f(*in_values)
if not isinstance(out_values, list):
out_values = [out_values]
# Check output types and shapes
assert len(out_dims) == len(out_values)
for d, value in zip(out_dims, out_values):
if scalar and d == 0:
# Should have been converted to a scalar value
assert isinstance(value, np.number)
else:
# Otherwise should be an array
assert isinstance(value, np.ndarray)
assert value.ndim == d
def test_aesara_function_bad_kwarg():
"""
Passing an unknown keyword argument to aesara_function() should raise an
exception.
"""
raises(Exception, lambda : aesara_function_([x], [x+1], foobar=3))
def test_slice():
assert aesara_code_(slice(1, 2, 3)) == slice(1, 2, 3)
def theq_slice(s1, s2):
for attr in ['start', 'stop', 'step']:
a1 = getattr(s1, attr)
a2 = getattr(s2, attr)
if a1 is None or a2 is None:
if not (a1 is None or a2 is None):
return False
elif not theq(a1, a2):
return False
return True
dtypes = {x: 'int32', y: 'int32'}
assert theq_slice(aesara_code_(slice(x, y), dtypes=dtypes), slice(xt, yt))
assert theq_slice(aesara_code_(slice(1, x, 3), dtypes=dtypes), slice(1, xt, 3))
def test_MatrixSlice():
cache = {}
n = sy.Symbol('n', integer=True)
X = sy.MatrixSymbol('X', n, n)
Y = X[1:2:3, 4:5:6]
Yt = aesara_code_(Y, cache=cache)
s = Scalar('int64')
assert tuple(Yt.owner.op.idx_list) == (slice(s, s, s), slice(s, s, s))
assert Yt.owner.inputs[0] == aesara_code_(X, cache=cache)
# == doesn't work in Aesara like it does in SymPy. You have to use
# equals.
assert all(Yt.owner.inputs[i].data == i for i in range(1, 7))
k = sy.Symbol('k')
aesara_code_(k, dtypes={k: 'int32'})
start, stop, step = 4, k, 2
Y = X[start:stop:step]
Yt = aesara_code_(Y, dtypes={n: 'int32', k: 'int32'})
# assert Yt.owner.op.idx_list[0].stop == kt
def test_BlockMatrix():
n = sy.Symbol('n', integer=True)
A, B, C, D = [sy.MatrixSymbol(name, n, n) for name in 'ABCD']
At, Bt, Ct, Dt = map(aesara_code_, (A, B, C, D))
Block = sy.BlockMatrix([[A, B], [C, D]])
Blockt = aesara_code_(Block)
solutions = [aet.join(0, aet.join(1, At, Bt), aet.join(1, Ct, Dt)),
aet.join(1, aet.join(0, At, Ct), aet.join(0, Bt, Dt))]
assert any(theq(Blockt, solution) for solution in solutions)
@SKIP
def test_BlockMatrix_Inverse_execution():
k, n = 2, 4
dtype = 'float32'
A = sy.MatrixSymbol('A', n, k)
B = sy.MatrixSymbol('B', n, n)
inputs = A, B
output = B.I*A
cutsizes = {A: [(n//2, n//2), (k//2, k//2)],
B: [(n//2, n//2), (n//2, n//2)]}
cutinputs = [sy.blockcut(i, *cutsizes[i]) for i in inputs]
cutoutput = output.subs(dict(zip(inputs, cutinputs)))
dtypes = dict(zip(inputs, [dtype]*len(inputs)))
f = aesara_function_(inputs, [output], dtypes=dtypes, cache={})
fblocked = aesara_function_(inputs, [sy.block_collapse(cutoutput)],
dtypes=dtypes, cache={})
ninputs = [np.random.rand(*x.shape).astype(dtype) for x in inputs]
ninputs = [np.arange(n*k).reshape(A.shape).astype(dtype),
np.eye(n).astype(dtype)]
ninputs[1] += np.ones(B.shape)*1e-5
assert np.allclose(f(*ninputs), fblocked(*ninputs), rtol=1e-5)
def test_DenseMatrix():
from aesara.tensor.basic import Join
t = sy.Symbol('theta')
for MatrixType in [sy.Matrix, sy.ImmutableMatrix]:
X = MatrixType([[sy.cos(t), -sy.sin(t)], [sy.sin(t), sy.cos(t)]])
tX = aesara_code_(X)
assert isinstance(tX, TensorVariable)
assert isinstance(tX.owner.op, Join)
def test_cache_basic():
""" Test single symbol-like objects are cached when printed by themselves. """
# Pairs of objects which should be considered equivalent with respect to caching
pairs = [
(x, sy.Symbol('x')),
(X, sy.MatrixSymbol('X', *X.shape)),
(f_t, sy.Function('f')(sy.Symbol('t'))),
]
for s1, s2 in pairs:
cache = {}
st = aesara_code_(s1, cache=cache)
# Test hit with same instance
assert aesara_code_(s1, cache=cache) is st
# Test miss with same instance but new cache
assert aesara_code_(s1, cache={}) is not st
# Test hit with different but equivalent instance
assert aesara_code_(s2, cache=cache) is st
def test_global_cache():
""" Test use of the global cache. """
from sympy.printing.aesaracode import global_cache
backup = dict(global_cache)
try:
# Temporarily empty global cache
global_cache.clear()
for s in [x, X, f_t]:
st = aesara_code(s)
assert aesara_code(s) is st
finally:
# Restore global cache
global_cache.update(backup)
def test_cache_types_distinct():
"""
Test that symbol-like objects of different types (Symbol, MatrixSymbol,
AppliedUndef) are distinguished by the cache even if they have the same
name.
"""
symbols = [sy.Symbol('f_t'), sy.MatrixSymbol('f_t', 4, 4), f_t]
cache = {} # Single shared cache
printed = {}
for s in symbols:
st = aesara_code_(s, cache=cache)
assert st not in printed.values()
printed[s] = st
# Check all printed objects are distinct
assert len(set(map(id, printed.values()))) == len(symbols)
# Check retrieving
for s, st in printed.items():
assert aesara_code(s, cache=cache) is st
def test_symbols_are_created_once():
"""
Test that a symbol is cached and reused when it appears in an expression
more than once.
"""
expr = sy.Add(x, x, evaluate=False)
comp = aesara_code_(expr)
assert theq(comp, xt + xt)
assert not theq(comp, xt + aesara_code_(x))
def test_cache_complex():
"""
Test caching on a complicated expression with multiple symbols appearing
multiple times.
"""
expr = x ** 2 + (y - sy.exp(x)) * sy.sin(z - x * y)
symbol_names = {s.name for s in expr.free_symbols}
expr_t = aesara_code_(expr)
# Iterate through variables in the Aesara computational graph that the
# printed expression depends on
seen = set()
for v in aesara.graph.basic.ancestors([expr_t]):
# Owner-less, non-constant variables should be our symbols
if v.owner is None and not isinstance(v, aesara.graph.basic.Constant):
# Check it corresponds to a symbol and appears only once
assert v.name in symbol_names
assert v.name not in seen
seen.add(v.name)
# Check all were present
assert seen == symbol_names
def test_Piecewise():
# A piecewise linear
expr = sy.Piecewise((0, x<0), (x, x<2), (1, True)) # ___/III
result = aesara_code_(expr)
assert result.owner.op == aet.switch
expected = aet.switch(xt<0, 0, aet.switch(xt<2, xt, 1))
assert theq(result, expected)
expr = sy.Piecewise((x, x < 0))
result = aesara_code_(expr)
expected = aet.switch(xt < 0, xt, np.nan)
assert theq(result, expected)
expr = sy.Piecewise((0, sy.And(x>0, x<2)), \
(x, sy.Or(x>2, x<0)))
result = aesara_code_(expr)
expected = aet.switch(aet.and_(xt>0,xt<2), 0, \
aet.switch(aet.or_(xt>2, xt<0), xt, np.nan))
assert theq(result, expected)
def test_Relationals():
assert theq(aesara_code_(sy.Eq(x, y)), aet.eq(xt, yt))
# assert theq(aesara_code_(sy.Ne(x, y)), aet.neq(xt, yt)) # TODO - implement
assert theq(aesara_code_(x > y), xt > yt)
assert theq(aesara_code_(x < y), xt < yt)
assert theq(aesara_code_(x >= y), xt >= yt)
assert theq(aesara_code_(x <= y), xt <= yt)
def test_complexfunctions():
dtypes = {x:'complex128', y:'complex128'}
xt, yt = aesara_code(x, dtypes=dtypes), aesara_code(y, dtypes=dtypes)
from sympy.functions.elementary.complexes import conjugate
from aesara.tensor import as_tensor_variable as atv
from aesara.tensor import complex as cplx
assert theq(aesara_code(y*conjugate(x), dtypes=dtypes), yt*(xt.conj()))
assert theq(aesara_code((1+2j)*x), xt*(atv(1.0)+atv(2.0)*cplx(0,1)))
def test_constantfunctions():
tf = aesara_function([],[1+1j])
assert(tf()==1+1j)