1import random
2from sympy import symbols, Derivative
3from sympy.tensor.array.expressions.array_expressions import ArrayTensorProduct, ArrayAdd, \
4    PermuteDims, ArrayDiagonal
5from sympy.core.relational import Eq, Ne, Ge, Gt, Le, Lt
6from sympy.external import import_module
7from sympy.functions import \
8    Abs, ceiling, exp, floor, sign, sin, asin, sqrt, cos, \
9    acos, tan, atan, atan2, cosh, acosh, sinh, asinh, tanh, atanh, \
10    re, im, arg, erf, loggamma, log
11from sympy.matrices import Matrix, MatrixBase, eye, randMatrix
12from sympy.matrices.expressions import \
13    Determinant, HadamardProduct, Inverse, MatrixSymbol, Trace
14from sympy.printing.tensorflow import tensorflow_code
15from sympy.tensor.array.expressions.conv_matrix_to_array import convert_matrix_to_array
16from sympy.utilities.lambdify import lambdify
17from sympy.testing.pytest import skip
18from sympy.testing.pytest import XFAIL
19
20
21tf = tensorflow = import_module("tensorflow")
22
23if tensorflow:
24    # Hide Tensorflow warnings
25    import os
26    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
27
28
29M = MatrixSymbol("M", 3, 3)
30N = MatrixSymbol("N", 3, 3)
31P = MatrixSymbol("P", 3, 3)
32Q = MatrixSymbol("Q", 3, 3)
33
34x, y, z, t = symbols("x y z t")
35
36if tf is not None:
37    llo = [[j for j in range(i, i+3)] for i in range(0, 9, 3)]
38    m3x3 = tf.constant(llo)
39    m3x3sympy = Matrix(llo)
40
41
42def _compare_tensorflow_matrix(variables, expr, use_float=False):
43    f = lambdify(variables, expr, 'tensorflow')
44    if not use_float:
45        random_matrices = [randMatrix(v.rows, v.cols) for v in variables]
46    else:
47        random_matrices = [randMatrix(v.rows, v.cols)/100. for v in variables]
48
49    graph = tf.Graph()
50    r = None
51    with graph.as_default():
52        random_variables = [eval(tensorflow_code(i)) for i in random_matrices]
53        session = tf.compat.v1.Session(graph=graph)
54        r = session.run(f(*random_variables))
55
56    e = expr.subs({k: v for k, v in zip(variables, random_matrices)})
57    e = e.doit()
58    if e.is_Matrix:
59        if not isinstance(e, MatrixBase):
60            e = e.as_explicit()
61        e = e.tolist()
62
63    if not use_float:
64        assert (r == e).all()
65    else:
66        r = [i for row in r for i in row]
67        e = [i for row in e for i in row]
68        assert all(
69            abs(a-b) < 10**-(4-int(log(abs(a), 10))) for a, b in zip(r, e))
70
71
72# Creating a custom inverse test.
73# See https://github.com/sympy/sympy/issues/18469
74def _compare_tensorflow_matrix_inverse(variables, expr, use_float=False):
75    f = lambdify(variables, expr, 'tensorflow')
76    if not use_float:
77        random_matrices = [eye(v.rows, v.cols)*4 for v in variables]
78    else:
79        random_matrices = [eye(v.rows, v.cols)*3.14 for v in variables]
80
81    graph = tf.Graph()
82    r = None
83    with graph.as_default():
84        random_variables = [eval(tensorflow_code(i)) for i in random_matrices]
85        session = tf.compat.v1.Session(graph=graph)
86        r = session.run(f(*random_variables))
87
88    e = expr.subs({k: v for k, v in zip(variables, random_matrices)})
89    e = e.doit()
90    if e.is_Matrix:
91        if not isinstance(e, MatrixBase):
92            e = e.as_explicit()
93        e = e.tolist()
94
95    if not use_float:
96        assert (r == e).all()
97    else:
98        r = [i for row in r for i in row]
99        e = [i for row in e for i in row]
100        assert all(
101            abs(a-b) < 10**-(4-int(log(abs(a), 10))) for a, b in zip(r, e))
102
103
104def _compare_tensorflow_matrix_scalar(variables, expr):
105    f = lambdify(variables, expr, 'tensorflow')
106    random_matrices = [
107        randMatrix(v.rows, v.cols).evalf() / 100 for v in variables]
108
109    graph = tf.Graph()
110    r = None
111    with graph.as_default():
112        random_variables = [eval(tensorflow_code(i)) for i in random_matrices]
113        session = tf.compat.v1.Session(graph=graph)
114        r = session.run(f(*random_variables))
115
116    e = expr.subs({k: v for k, v in zip(variables, random_matrices)})
117    e = e.doit()
118    assert abs(r-e) < 10**-6
119
120
121def _compare_tensorflow_scalar(
122    variables, expr, rng=lambda: random.randint(0, 10)):
123    f = lambdify(variables, expr, 'tensorflow')
124    rvs = [rng() for v in variables]
125
126    graph = tf.Graph()
127    r = None
128    with graph.as_default():
129        tf_rvs = [eval(tensorflow_code(i)) for i in rvs]
130        session = tf.compat.v1.Session(graph=graph)
131        r = session.run(f(*tf_rvs))
132
133    e = expr.subs({k: v for k, v in zip(variables, rvs)}).evalf().doit()
134    assert abs(r-e) < 10**-6
135
136
137def _compare_tensorflow_relational(
138    variables, expr, rng=lambda: random.randint(0, 10)):
139    f = lambdify(variables, expr, 'tensorflow')
140    rvs = [rng() for v in variables]
141
142    graph = tf.Graph()
143    r = None
144    with graph.as_default():
145        tf_rvs = [eval(tensorflow_code(i)) for i in rvs]
146        session = tf.compat.v1.Session(graph=graph)
147        r = session.run(f(*tf_rvs))
148
149    e = expr.subs({k: v for k, v in zip(variables, rvs)}).doit()
150    assert r == e
151
152
153def test_tensorflow_printing():
154    assert tensorflow_code(eye(3)) == \
155        "tensorflow.constant([[1, 0, 0], [0, 1, 0], [0, 0, 1]])"
156
157    expr = Matrix([[x, sin(y)], [exp(z), -t]])
158    assert tensorflow_code(expr) == \
159        "tensorflow.Variable(" \
160            "[[x, tensorflow.math.sin(y)]," \
161            " [tensorflow.math.exp(z), -t]])"
162
163
164# This (random) test is XFAIL because it fails occasionally
165# See https://github.com/sympy/sympy/issues/18469
166@XFAIL
167def test_tensorflow_math():
168    if not tf:
169        skip("TensorFlow not installed")
170
171    expr = Abs(x)
172    assert tensorflow_code(expr) == "tensorflow.math.abs(x)"
173    _compare_tensorflow_scalar((x,), expr)
174
175    expr = sign(x)
176    assert tensorflow_code(expr) == "tensorflow.math.sign(x)"
177    _compare_tensorflow_scalar((x,), expr)
178
179    expr = ceiling(x)
180    assert tensorflow_code(expr) == "tensorflow.math.ceil(x)"
181    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
182
183    expr = floor(x)
184    assert tensorflow_code(expr) == "tensorflow.math.floor(x)"
185    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
186
187    expr = exp(x)
188    assert tensorflow_code(expr) == "tensorflow.math.exp(x)"
189    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
190
191    expr = sqrt(x)
192    assert tensorflow_code(expr) == "tensorflow.math.sqrt(x)"
193    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
194
195    expr = x ** 4
196    assert tensorflow_code(expr) == "tensorflow.math.pow(x, 4)"
197    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
198
199    expr = cos(x)
200    assert tensorflow_code(expr) == "tensorflow.math.cos(x)"
201    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
202
203    expr = acos(x)
204    assert tensorflow_code(expr) == "tensorflow.math.acos(x)"
205    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.uniform(0, 0.95))
206
207    expr = sin(x)
208    assert tensorflow_code(expr) == "tensorflow.math.sin(x)"
209    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
210
211    expr = asin(x)
212    assert tensorflow_code(expr) == "tensorflow.math.asin(x)"
213    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
214
215    expr = tan(x)
216    assert tensorflow_code(expr) == "tensorflow.math.tan(x)"
217    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
218
219    expr = atan(x)
220    assert tensorflow_code(expr) == "tensorflow.math.atan(x)"
221    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
222
223    expr = atan2(y, x)
224    assert tensorflow_code(expr) == "tensorflow.math.atan2(y, x)"
225    _compare_tensorflow_scalar((y, x), expr, rng=lambda: random.random())
226
227    expr = cosh(x)
228    assert tensorflow_code(expr) == "tensorflow.math.cosh(x)"
229    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.random())
230
231    expr = acosh(x)
232    assert tensorflow_code(expr) == "tensorflow.math.acosh(x)"
233    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.uniform(1, 2))
234
235    expr = sinh(x)
236    assert tensorflow_code(expr) == "tensorflow.math.sinh(x)"
237    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.uniform(1, 2))
238
239    expr = asinh(x)
240    assert tensorflow_code(expr) == "tensorflow.math.asinh(x)"
241    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.uniform(1, 2))
242
243    expr = tanh(x)
244    assert tensorflow_code(expr) == "tensorflow.math.tanh(x)"
245    _compare_tensorflow_scalar((x,), expr, rng=lambda: random.uniform(1, 2))
246
247    expr = atanh(x)
248    assert tensorflow_code(expr) == "tensorflow.math.atanh(x)"
249    _compare_tensorflow_scalar(
250        (x,), expr, rng=lambda: random.uniform(-.5, .5))
251
252    expr = erf(x)
253    assert tensorflow_code(expr) == "tensorflow.math.erf(x)"
254    _compare_tensorflow_scalar(
255        (x,), expr, rng=lambda: random.random())
256
257    expr = loggamma(x)
258    assert tensorflow_code(expr) == "tensorflow.math.lgamma(x)"
259    _compare_tensorflow_scalar(
260        (x,), expr, rng=lambda: random.random())
261
262
263def test_tensorflow_complexes():
264    assert tensorflow_code(re(x)) == "tensorflow.math.real(x)"
265    assert tensorflow_code(im(x)) == "tensorflow.math.imag(x)"
266    assert tensorflow_code(arg(x)) == "tensorflow.math.angle(x)"
267
268
269def test_tensorflow_relational():
270    if not tf:
271        skip("TensorFlow not installed")
272
273    expr = Eq(x, y)
274    assert tensorflow_code(expr) == "tensorflow.math.equal(x, y)"
275    _compare_tensorflow_relational((x, y), expr)
276
277    expr = Ne(x, y)
278    assert tensorflow_code(expr) == "tensorflow.math.not_equal(x, y)"
279    _compare_tensorflow_relational((x, y), expr)
280
281    expr = Ge(x, y)
282    assert tensorflow_code(expr) == "tensorflow.math.greater_equal(x, y)"
283    _compare_tensorflow_relational((x, y), expr)
284
285    expr = Gt(x, y)
286    assert tensorflow_code(expr) == "tensorflow.math.greater(x, y)"
287    _compare_tensorflow_relational((x, y), expr)
288
289    expr = Le(x, y)
290    assert tensorflow_code(expr) == "tensorflow.math.less_equal(x, y)"
291    _compare_tensorflow_relational((x, y), expr)
292
293    expr = Lt(x, y)
294    assert tensorflow_code(expr) == "tensorflow.math.less(x, y)"
295    _compare_tensorflow_relational((x, y), expr)
296
297
298# This (random) test is XFAIL because it fails occasionally
299# See https://github.com/sympy/sympy/issues/18469
300@XFAIL
301def test_tensorflow_matrices():
302    if not tf:
303        skip("TensorFlow not installed")
304
305    expr = M
306    assert tensorflow_code(expr) == "M"
307    _compare_tensorflow_matrix((M,), expr)
308
309    expr = M + N
310    assert tensorflow_code(expr) == "tensorflow.math.add(M, N)"
311    _compare_tensorflow_matrix((M, N), expr)
312
313    expr = M * N
314    assert tensorflow_code(expr) == "tensorflow.linalg.matmul(M, N)"
315    _compare_tensorflow_matrix((M, N), expr)
316
317    expr = HadamardProduct(M, N)
318    assert tensorflow_code(expr) == "tensorflow.math.multiply(M, N)"
319    _compare_tensorflow_matrix((M, N), expr)
320
321    expr = M*N*P*Q
322    assert tensorflow_code(expr) == \
323        "tensorflow.linalg.matmul(" \
324            "tensorflow.linalg.matmul(" \
325                "tensorflow.linalg.matmul(M, N), P), Q)"
326    _compare_tensorflow_matrix((M, N, P, Q), expr)
327
328    expr = M**3
329    assert tensorflow_code(expr) == \
330        "tensorflow.linalg.matmul(tensorflow.linalg.matmul(M, M), M)"
331    _compare_tensorflow_matrix((M,), expr)
332
333    expr = Trace(M)
334    assert tensorflow_code(expr) == "tensorflow.linalg.trace(M)"
335    _compare_tensorflow_matrix((M,), expr)
336
337    expr = Determinant(M)
338    assert tensorflow_code(expr) == "tensorflow.linalg.det(M)"
339    _compare_tensorflow_matrix_scalar((M,), expr)
340
341    expr = Inverse(M)
342    assert tensorflow_code(expr) == "tensorflow.linalg.inv(M)"
343    _compare_tensorflow_matrix_inverse((M,), expr, use_float=True)
344
345    expr = M.T
346    assert tensorflow_code(expr, tensorflow_version='1.14') == \
347        "tensorflow.linalg.matrix_transpose(M)"
348    assert tensorflow_code(expr, tensorflow_version='1.13') == \
349        "tensorflow.matrix_transpose(M)"
350
351    _compare_tensorflow_matrix((M,), expr)
352
353
354def test_codegen_einsum():
355    if not tf:
356        skip("TensorFlow not installed")
357
358    graph = tf.Graph()
359    with graph.as_default():
360        session = tf.compat.v1.Session(graph=graph)
361
362        M = MatrixSymbol("M", 2, 2)
363        N = MatrixSymbol("N", 2, 2)
364
365        cg = convert_matrix_to_array(M * N)
366        f = lambdify((M, N), cg, 'tensorflow')
367
368        ma = tf.constant([[1, 2], [3, 4]])
369        mb = tf.constant([[1,-2], [-1, 3]])
370        y = session.run(f(ma, mb))
371        c = session.run(tf.matmul(ma, mb))
372        assert (y == c).all()
373
374
375def test_codegen_extra():
376    if not tf:
377        skip("TensorFlow not installed")
378
379    graph = tf.Graph()
380    with graph.as_default():
381        session = tf.compat.v1.Session()
382
383        M = MatrixSymbol("M", 2, 2)
384        N = MatrixSymbol("N", 2, 2)
385        P = MatrixSymbol("P", 2, 2)
386        Q = MatrixSymbol("Q", 2, 2)
387        ma = tf.constant([[1, 2], [3, 4]])
388        mb = tf.constant([[1,-2], [-1, 3]])
389        mc = tf.constant([[2, 0], [1, 2]])
390        md = tf.constant([[1,-1], [4, 7]])
391
392        cg = ArrayTensorProduct(M, N)
393        assert tensorflow_code(cg) == \
394            'tensorflow.linalg.einsum("ab,cd", M, N)'
395        f = lambdify((M, N), cg, 'tensorflow')
396        y = session.run(f(ma, mb))
397        c = session.run(tf.einsum("ij,kl", ma, mb))
398        assert (y == c).all()
399
400        cg = ArrayAdd(M, N)
401        assert tensorflow_code(cg) == 'tensorflow.math.add(M, N)'
402        f = lambdify((M, N), cg, 'tensorflow')
403        y = session.run(f(ma, mb))
404        c = session.run(ma + mb)
405        assert (y == c).all()
406
407        cg = ArrayAdd(M, N, P)
408        assert tensorflow_code(cg) == \
409            'tensorflow.math.add(tensorflow.math.add(M, N), P)'
410        f = lambdify((M, N, P), cg, 'tensorflow')
411        y = session.run(f(ma, mb, mc))
412        c = session.run(ma + mb + mc)
413        assert (y == c).all()
414
415        cg = ArrayAdd(M, N, P, Q)
416        assert tensorflow_code(cg) == \
417            'tensorflow.math.add(' \
418                'tensorflow.math.add(tensorflow.math.add(M, N), P), Q)'
419        f = lambdify((M, N, P, Q), cg, 'tensorflow')
420        y = session.run(f(ma, mb, mc, md))
421        c = session.run(ma + mb + mc + md)
422        assert (y == c).all()
423
424        cg = PermuteDims(M, [1, 0])
425        assert tensorflow_code(cg) == 'tensorflow.transpose(M, [1, 0])'
426        f = lambdify((M,), cg, 'tensorflow')
427        y = session.run(f(ma))
428        c = session.run(tf.transpose(ma))
429        assert (y == c).all()
430
431        cg = PermuteDims(ArrayTensorProduct(M, N), [1, 2, 3, 0])
432        assert tensorflow_code(cg) == \
433            'tensorflow.transpose(' \
434                'tensorflow.linalg.einsum("ab,cd", M, N), [1, 2, 3, 0])'
435        f = lambdify((M, N), cg, 'tensorflow')
436        y = session.run(f(ma, mb))
437        c = session.run(tf.transpose(tf.einsum("ab,cd", ma, mb), [1, 2, 3, 0]))
438        assert (y == c).all()
439
440        cg = ArrayDiagonal(ArrayTensorProduct(M, N), (1, 2))
441        assert tensorflow_code(cg) == \
442            'tensorflow.linalg.einsum("ab,bc->acb", M, N)'
443        f = lambdify((M, N), cg, 'tensorflow')
444        y = session.run(f(ma, mb))
445        c = session.run(tf.einsum("ab,bc->acb", ma, mb))
446        assert (y == c).all()
447
448
449def test_MatrixElement_printing():
450    A = MatrixSymbol("A", 1, 3)
451    B = MatrixSymbol("B", 1, 3)
452    C = MatrixSymbol("C", 1, 3)
453
454    assert tensorflow_code(A[0, 0]) == "A[0, 0]"
455    assert tensorflow_code(3 * A[0, 0]) == "3*A[0, 0]"
456
457    F = C[0, 0].subs(C, A - B)
458    assert tensorflow_code(F) == "(tensorflow.math.add((-1)*B, A))[0, 0]"
459
460
461def test_tensorflow_Derivative():
462    expr = Derivative(sin(x), x)
463    assert tensorflow_code(expr) == \
464        "tensorflow.gradients(tensorflow.math.sin(x), x)[0]"
465