def testCaseSimple(self): with self.session() as sess: def my_graph(pa, pb, pc): with ipu.scopes.ipu_scope("/device:IPU:0"): @eager_function.defun def b0(x, y): return x + y @eager_function.defun def b1(x, y): return x - y @eager_function.defun def b2(x, y): return x * y branches = [ f.get_concrete_function(array_ops.zeros_like(pb), array_ops.zeros_like(pc)) for f in [b0, b1, b2] ] c_out = gen_functional_ops.case(pa, input=[pb, pc], Tout=[dtypes.float32], branches=branches) return [c_out[0]] with ops.device('cpu'): pa = array_ops.placeholder(np.int32, [], name="a") pb = array_ops.placeholder(np.float32, [2], name="b") pc = array_ops.placeholder(np.float32, [2], name="c") out = ipu.ipu_compiler.compile(my_graph, [pa, pb, pc]) report = ReportJSON(self, sess) report.reset() result = sess.run(out, {pa: 0, pb: [0., 1.], pc: [1., 5.]}) self.assertAllClose(result[0], [1., 6.]) result = sess.run(out, {pa: 1, pb: [0., 1.], pc: [1., 5.]}) self.assertAllClose(result[0], [-1., -4.]) result = sess.run(out, {pa: 2, pb: [0., 1.], pc: [1., 5.]}) self.assertAllClose(result[0], [0., 5.]) result = sess.run(out, {pa: 10, pb: [0., 1.], pc: [1., 5.]}) self.assertAllClose(result[0], [0., 5.]) report.parse_log() report.assert_contains_one_compile_event()
def testIpuWhilePerfTest(self): with self.session() as sess: def cond(i, v): del v return math_ops.less(i, 15) def body(i, v): v = v + i i = i + 1 return (i, v) def my_net(v): i = constant_op.constant(0) r = control_flow_ops.while_loop(cond, body, (i, v), maximum_iterations=10) return [r[1]] with ops.device('cpu'): v = array_ops.placeholder(np.int32, [500]) with ipu.scopes.ipu_scope("/device:IPU:0"): r = ipu.ipu_compiler.compile(my_net, inputs=[v]) report = ReportJSON(self, sess) report.reset() result = sess.run(r, {v: np.zeros([500], np.int32)}) self.assertAllClose(result[0], np.broadcast_to(45, [500])) report.parse_log() # Check that there is only one real compile report.assert_contains_one_compile_event() # Check that there is only one execute report.assert_num_execution_reports_equal(1)