do it

Author: MarisaKirisame

tests/python/relay/test_pass_gradient.py

@@ -156,13 +156,20 @@ def test_tuple():
 def test_pow():
     mod = relay.Module()
     p = Prelude(mod)
-    shape = ()
+    shape = (10, 10)
     dtype = 'float32'
     t = relay.TensorType(shape, dtype)
     x = relay.var("x", t)
     double = relay.Function([x], x + x)
     i = relay.var("i", t)
-    func = relay.Function([x], p.iterate)
+    func = relay.Function([i], relay.Call(p.iterate(double, p.s(p.s(p.s(p.z())))), [i]))
+    back_func = relay.ir_pass.infer_type(gradient(func, mod=mod), mod=mod)
+    assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])]))
+    i_nd = rand(dtype, *shape)
+    ex = create_executor(mod=mod)
+    forward, (grad_i,) = ex.evaluate(back_func)(i_nd)
+    np.testing.assert_allclose(forward.asnumpy(), 8 * i_nd.asnumpy())
+    np.testing.assert_allclose(grad_i.asnumpy(), 8 * np.ones_like(grad_i.asnumpy()))
 
 if __name__ == "__main__":
     test_id()