def test_programme():
engine = MyEngine()
code = CodeSegment(engine)
code.batch(u='a', v='d')
code.batch_with_exarg(u='a', v='e', factor=3.0)
(d, e), tape = code.compute(('d', 'e'), {'a' : 1.0}, return_tape=True)
assert_array_equal(d, 2.0)
assert_array_equal(e, 6.0)
e, d, _a = code.compute_with_gradient(['e', 'd', '_a'], {'a' : 1.0}, {'_d': 1.0, '_e' : 0.0})
assert_array_equal(d, 2.0)
assert_array_equal(e, 6.0)
assert_array_equal(_a, 2.0)
e, d, _a = code.compute_with_gradient(['e', 'd', '_a'], {'a' : 1.0}, {'_d': 0.0, '_e' : 1.0})
assert_array_equal(d, 2.0)
assert_array_equal(e, 6.0)
assert_array_equal(_a, 6.0)
def test_programme_unused():
engine = MyEngine()
code = CodeSegment(engine)
code.batch_unused(u='a', v='d')
d, _a = code.compute_with_gradient(['d', '_a'], {'a' : 1.0}, {'_d': 1.0})
assert_array_equal(d, 1.0)
assert_array_equal(_a, 0.0)
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_zeros():
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)
a, b1, b2, _a = code.compute_with_gradient(['a', 'b1', 'b2', '_a'], {'a' : 1.0}, {'_b1': ZERO, '_b2' : ZERO, '_a' : ZERO})
assert _a is ZERO
def test_nested_compute():
engine = MyEngine()
code = CodeSegment(engine)
code.unitary(x='a', y='b', factor=3.0)
code.unitary(x='b', y='c', factor=3.0)
c = code.compute('c', {'a' : 1.0})
assert_array_equal(c, 9.0)
c, _a = code.compute_with_gradient(['c', '_a'], {'a' : 1.0}, {'_c': 1.0})
assert_array_equal(c, 9.0)
assert_array_equal(_a, 9.0)
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_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)
