def test_literal(): engine = MyEngine() code = CodeSegment(engine) code.binary(x1='a', x2=Literal(2.0), y='a') code.batch(u=Literal(2.0), v='d') d, a, _a = code.compute_with_gradient(['d', 'a', '_a'], {'a' : 1.0}, {'_a': 1.0, '_d' : 0.0}) assert_array_equal(a, 3.0) assert_array_equal(_a, 1.0)
def test_copy(): engine = MyEngine() code = CodeSegment(engine) code.batch_with_sub(u='a', v='e') code.unitary(x='a', y='a', factor=3.0) code.unitary(x='a', y='b1', factor=3.0) code.unitary(x='a', y='b2', factor=3.0) code.binary(x1='b1', x2='b2', y='b1') code.unitary(x='b1', y='d', factor=3.0) code.batch(u='b2', v='f') code2 = code.copy()
def test_to_graph(): engine = MyEngine() code = CodeSegment(engine) code.batch_with_sub(u='a', v='e') code.unitary(x='a', y='a', factor=3.0) code.unitary(x='a', y='b1', factor=3.0) code.unitary(x='a', y='b2', factor=3.0) code.binary(x1='b1', x2='b2', y='b1') code.unitary(x='b1', y='d', factor=3.0) code.batch(u='b2', v='f') d, tape = code.compute(('e', 'a', 'f', 'd'), {'a' : 1.0}, return_tape=True) vjp = tape.get_vjp() graph1 = code.to_graph() graph2 = vjp.to_graph()
def test_inplace(): engine = MyEngine() code = CodeSegment(engine) code.unitary(x='a', y='a', factor=3.0) code.unitary(x='a', y='b1', factor=3.0) code.unitary(x='a', y='b2', factor=3.0) code.binary(x1='b1', x2='b2', y='b1') code.unitary(x='b1', y='d', factor=3.0) d = code.compute('d', {'a' : 1.0}) assert_array_equal(d, 54.0) d, _a = code.compute_with_gradient(['d', '_a'], {'a' : 1.0}, {'_d': 1.0}) assert_array_equal(d, 54.0) assert_array_equal(_a, 54.0)
def test_jvp(): engine = MyEngine() code = CodeSegment(engine) code.unitary(x='a', y='b1', factor=3.0) code.unitary(x='a', y='b2', factor=3.0) code.unitary(x='a', y='b3', factor=3.0) code.unitary(x='a', y='b4', factor=3.0) code.binary(x1='b1', x2='b2', y='c1') code.binary(x1='b3', x2='b4', y='c2') code.binary(x1='c1', x2='c2', y='d') jvp = code.get_jvp(init={'a' : 1.0}) d_ = jvp.compute('d_', {'a_' : 1.0}) assert_array_equal(d_, 12.0) c1_, d_ = jvp.compute(['c1_', 'd_'], {'a_' : 1.0}) assert_array_equal(d_, 12.0) assert_array_equal(c1_, 6.0)
def test_tape_gradients(): engine = MyEngine() code = CodeSegment(engine) code.unitary(x='a', y='b1', factor=3.0) code.unitary(x='a', y='b2', factor=3.0) code.unitary(x='a', y='b3', factor=3.0) code.unitary(x='a', y='b4', factor=3.0) code.binary(x1='b1', x2='b2', y='c1') code.binary(x1='b3', x2='b4', y='c2') code.binary(x1='c1', x2='c2', y='d') d, tape = code.compute('d', {'a' : 1.0}, return_tape=True) assert_array_equal(d, 12.0) vjp = tape.get_vjp() _a = vjp.compute(['_a'], {'_d': 1.0}) assert_array_equal(_a, 12.0) jvp = tape.get_jvp() d_ = jvp.compute(['d_'], {'a_': 1.0}) assert_array_equal(d_, 12.0) d, _a = code.compute_with_gradient(['d', '_a'], {'a' : 1.0}, {'_d': 1.0}) assert_array_equal(d, 12.0) assert_array_equal(_a, 12.0) (c1, d), tape = code.compute(['c1', 'd'], {'a' : 1.0}, return_tape=True) assert_array_equal(d, 12.0) assert_array_equal(c1, 6.0) vjp = tape.get_vjp() jvp = tape.get_jvp() _a = vjp.compute(['_a'], {'_c1': 1.0}) c1_, d_ = jvp.compute(['c1_', 'd_'], {'a_': 1.0}) assert_array_equal(c1_, 6.0) assert_array_equal(d_, 12.0)
def batch_with_exarg(engine, u, v, factor): code = CodeSegment(engine) code.unitary(x=u, y=u, factor=factor) code.binary(x1=u, x2=u, y=v) return code
def batch(engine, u, v): code = CodeSegment(engine) code.unitary(x=u, y=u, factor=1.0) code.binary(x1=u, x2=u, y=v) return code