def test_derivative():
m = TensorProductMean(lambda x: x**2)
m2 = TensorProductMean(lambda x: x**3)
x = B.randn(tf.float64, 10, 1)
allclose(m.diff(0)(x), 2 * x)
allclose(m2.diff(0)(x), 3 * x**2)
def test_derivative():
B.backend_to_tf()
s = B.Session()
m = TensorProductMean(lambda x: x**2)
m2 = TensorProductMean(lambda x: x**3)
x = B.array(np.random.randn(10, 1))
yield assert_allclose, s.run(m.diff(0)(x)), s.run(2 * x)
yield assert_allclose, s.run(m2.diff(0)(x)), s.run(3 * x**2)
s.close()
B.backend_to_np()