def test_jvp_programme(): engine = MyEngine() code = CodeSegment(engine) code.batch(u='a', v='d') code.unitary(x='d', y='d', factor=2.0) code.batch_with_exarg(u='a', v='e', factor=3.0) jvp = code.get_jvp(init={'a' : 1.0}) d_, e_ = jvp.compute(['d_', 'e_'], {'a_' : 1.0}) assert_array_equal(d_, 4.0) assert_array_equal(e_, 6.0)
def test_jvp_vector(): engine = MyEngine() code = CodeSegment(engine) code.batch(u='a', v='d') code.unitary(x='d', y='d', factor=2.0) code.batch_with_exarg(u='a', v='e', factor=3.0) A = numpy.array jvp = code.get_jvp(init={'a' : A([1.0, 1.0])}) d_, e_ = jvp.compute(['d_', 'e_'], {'a_' : A([1.0, 1.0])}) assert_array_equal(d_, [4.0, 4.0]) assert_array_equal(e_, [6.0, 6.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 batch_batch(engine, u, v): code = CodeSegment(engine) code.batch_with_exarg(u=u, v='t', factor=1) code.unitary(x='t', y=v, factor=1.0) return code