def testCommonSubexpressionElimination(self):
def f1(x):
return 3 * x**2 + x**2
def f2(x):
y = x**2
return 3 * y + y
expr1 = tracers.make_expr(f1, 1)
expr2 = tracers.make_expr(f2, 1)
code1 = tracers.print_expr(expr1)
code2 = tracers.print_expr(expr2)
self.assertGreater(len(code1), len(code2))
code1_cse = tracers.print_expr(tracers.remake_expr(expr1)) # applies cse
self.assertEqual(len(code1_cse), len(code2))
def testPrintExpr(self):
def fun(x, y):
return 2 * x**2 + np.tanh(y)
expr = tracers.make_expr(fun, 4, 5)
printed_expr = tracers.print_expr(expr)
expected = ("temp_0 = power(x, 2)\n"
"temp_1 = multiply(2, temp_0)\n"
"temp_2 = tanh(y)\n"
"temp_3 = add(temp_1, temp_2)\n")
self.assertEqual(printed_expr, expected)
def testEinsumTransposeRewriter(self):
def fun(x, y):
return np.einsum('ij,j->i', x.T, y)
x = npr.randn(4, 3)
y = npr.randn(4)
expr = tracers.make_expr(fun, x, y)
self.assertEqual(expr.expr_node.fun.__name__, 'einsum')
expr = self._rewriter_test_helper(fun,
rewrites.transpose_inside_einsum, x,
y)
self.assertFalse('transpose' in tracers.print_expr(expr))
def testInlineExpr(self):
def f(x, y):
return 2 * x + y
def g(z):
return 3 * z + z**2
expr = tracers.make_expr(f, 1, 2)
subexpr = tracers.make_expr(g, 3)
target_node = expr.expr_node.parents[0]
new_expr = tracers.inline_expr(subexpr, {'z': target_node})
printed_expr = tracers.print_expr(new_expr)
expected = ("temp_0 = multiply(2, x)\n"
"temp_1 = multiply(3, temp_0)\n"
"temp_2 = power(temp_0, 2)\n"
"temp_3 = add(temp_1, temp_2)\n")
self.assertEqual(printed_expr, expected)
self.assertEqual(3 * (2 * 5) + (2 * 5)**2,
tracers.eval_expr(new_expr, {'x': 5}))
self.assertEqual(f(6, 7), tracers.eval_expr(expr, {'x': 6, 'y': 7}))