def test_vector():
engine = VectorEngine()
code = CodeSegment(engine)
code.multiply(x='x1', y='r1', factor=3.0)
code.multiply(x='x2', y='r2', factor=3.0)
code.dot(x1='r1', x2='r2', y='r3')
code.power(x='r3', y='y', factor=2)
init = {'x1': numpy.array([1.0, 1.0]), 'x2': numpy.array([1.0, 1.0])}
b, tape = code.compute('y', init, return_tape=True)
backward_gradient = tape.get_vjp()
forward_gradient = code.get_jvp()
_x1, _x2 = backward_gradient.compute(['_x1', '_x2'], {'_y': 1.0})
for (x1_, x2_), ind in bases((2, 2), return_index=True):
d = {'x1_': x1_, 'x2_': x2_}
d.update(init)
y_ = forward_gradient.compute('y_', d)
assert_array_equal(y_, numpy.array([_x1, _x2])[ind])
x1_, x2_ = impulse((2, 2), (0, 1))
d = {'x1_': x1_, 'x2_': x2_}
d.update(init)
(y, y_), tape = forward_gradient.compute(['y', 'y_'], d, return_tape=True)
hessian_dot = tape.get_vjp()
_x1, _x2 = hessian_dot.compute({'_x1', '_x2'}, {'_y_': 1.0, '_y': 0.0})
def test_optimized_execution():
engine = MyEngine()
code = CodeSegment(engine)
code.unitary(x='a', y='d', factor=3.0)
code.unitary(x='a', y='b', factor=3.0)
code.unitary(x='a', y='c', factor=3.0)
opt = code.optimize(['b'])
assert len(opt.nodes) == 1
b, tape = code.compute('b', {'a' : 1.0}, return_tape=True)
print(tape)
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 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_scalar():
engine = ScalarEngine()
code = CodeSegment(engine)
code.unitary(x='a', y='r', factor=9.0)
code.power(x='r', y='b', factor=2.0)
b = code.compute('b', {'a': 2.0})
b, tape = code.compute('b', {'a': 2.0}, return_tape=True)
backward_gradient = tape.get_vjp()
forward_gradient = code.get_jvp()
b_ = forward_gradient.compute('b_', {'a': 2.0, 'a_': 1.0})
_a = backward_gradient.compute('_a', {'_b': 1.0})
assert_array_equal(b_, _a)
# to do the hessian, first augment the compute with a forward pass
# then do the backward gradient on the forward tape
# we need us terms like x_, etc.
print(forward_gradient)
(b_, b), tape = forward_gradient.compute(['b_', 'b'], {
'a': 2.0,
'a_': 1.0
},
return_tape=True)
print('b_', 'b', b_, b)
hessian_dot = tape.get_vjp()
print(hessian_dot)
_a = hessian_dot.compute(
'_a', {
'_b_': 1.0,
'_b': 0.0
},
monitor=lambda node, frontier, r: print('---', node, frontier, r))
print(_a)
def test_jvp_programme_nested():
engine = MyEngine()
code = CodeSegment(engine)
code.batch_batch(u='a', v='d')
jvp = code.get_jvp(init={'a' : 1.0})
d_ = jvp.compute('d_', {'a_' : 1.0})
assert_array_equal(d_, 2.0)
d, tape = code.compute('d', init={'a' : 1.0}, return_tape=True)
jvp = tape.get_jvp()
d_ = jvp.compute('d_', {'a_' : 1.0})
assert_array_equal(d_, 2.0)
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_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_compute():
engine = MyEngine()
code = CodeSegment(engine)
code.unitary(x='a', y='b', factor=3.0)
b = code.compute('b', {'a' : 1.0})
assert_array_equal(b, 3.0)
