def test_tuple(): shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) y = relay.var("y", t) z = relay.var("z", t) tup = relay.Var("tup") func = relay.Function( [x, y, z], relay.Let( tup, relay.Tuple([x, y, z]), relay.TupleGetItem(tup, 0) + relay.TupleGetItem(tup, 1) - relay.TupleGetItem(tup, 2))) back_func = relay.ir_pass.infer_type(gradient(func)) assert back_func.checked_type == relay.FuncType( [t, t, t], relay.TupleType([t, relay.TupleType([t, t, t])])) x_nd = rand(dtype, *shape) y_nd = rand(dtype, *shape) z_nd = rand(dtype, *shape) x_np = x_nd.asnumpy() y_np = y_nd.asnumpy() z_np = z_nd.asnumpy() expected_forward = x_np + y_np - z_np ex = create_executor() forward, (grad_x, grad_y, grad_z) = ex.evaluate(back_func)(x_nd, y_nd, z_nd) tvm.testing.assert_allclose(forward.asnumpy(), expected_forward) tvm.testing.assert_allclose(grad_x.asnumpy(), np.ones_like(grad_x.asnumpy())) tvm.testing.assert_allclose(grad_y.asnumpy(), np.ones_like(grad_y.asnumpy())) tvm.testing.assert_allclose(grad_z.asnumpy(), -1 * np.ones_like(grad_z.asnumpy()))
def test_broadcast_add(): shape1 = (3, 4, 1) shape2 = (1, 5) dtype = 'float32' x_nd = rand(dtype, *shape1) y_nd = rand(dtype, *shape2) x_np = x_nd.asnumpy() y_np = y_nd.asnumpy() expected_forward = x_np + y_np t1 = relay.TensorType(shape1, dtype) t2 = relay.TensorType(shape2, dtype) x = relay.var("x", t1) y = relay.var("y", t2) func = relay.Function([x, y], x + y) full_func = relay.ir_pass.infer_type(gradient(func)) assert full_func.checked_type == relay.FuncType([t1, t2], relay.TupleType([relay.TensorType(expected_forward.shape, dtype), relay.TupleType([t1, t2])])) ex = create_executor() forward, (grad_x, grad_y) = ex.evaluate(full_func)(x_nd, y_nd) tvm.testing.assert_allclose(forward.asnumpy(), expected_forward) tvm.testing.assert_allclose(grad_x.asnumpy(), np.ones_like(expected_forward).sum(axis=2, keepdims=True)) tvm.testing.assert_allclose(grad_y.asnumpy(), np.ones_like(expected_forward).sum(axis=(0, 1), keepdims=True).squeeze(axis=0))
def test_broadcast_add(): shape1 = (3, 4, 1) shape2 = (1, 5) dtype = 'float32' x_nd = rand(dtype, *shape1) y_nd = rand(dtype, *shape2) x_np = x_nd.asnumpy() y_np = y_nd.asnumpy() expected_forward = x_np + y_np t1 = relay.TensorType(shape1, dtype) t2 = relay.TensorType(shape2, dtype) x = relay.var("x", t1) y = relay.var("y", t2) func = relay.Function([x, y], x + y) full_func = relay.ir_pass.infer_type(gradient(func)) assert full_func.checked_type == relay.FuncType( [t1, t2], relay.TupleType([ relay.TensorType(expected_forward.shape, dtype), relay.TupleType([t1, t2]) ])) ex = create_executor() forward, (grad_x, grad_y) = ex.evaluate(full_func)(x_nd, y_nd) tvm.testing.assert_allclose(forward.asnumpy(), expected_forward) tvm.testing.assert_allclose( grad_x.asnumpy(), np.ones_like(expected_forward).sum(axis=2, keepdims=True)) tvm.testing.assert_allclose( grad_y.asnumpy(), np.ones_like(expected_forward).sum(axis=(0, 1), keepdims=True).squeeze(axis=0))
def test_tuple(): shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) y = relay.var("y", t) z = relay.var("z", t) tup = relay.Var("tup") func = relay.Function([x, y, z], relay.Let(tup, relay.Tuple([x, y, z]), relay.TupleGetItem(tup, 0) + relay.TupleGetItem(tup, 1) - relay.TupleGetItem(tup, 2))) back_func = relay.ir_pass.infer_type(gradient(func)) assert back_func.checked_type == relay.FuncType([t, t, t], relay.TupleType([t, relay.TupleType([t, t, t])])) x_nd = rand(dtype, *shape) y_nd = rand(dtype, *shape) z_nd = rand(dtype, *shape) x_np = x_nd.asnumpy() y_np = y_nd.asnumpy() z_np = z_nd.asnumpy() expected_forward = x_np + y_np - z_np ex = create_executor() forward, (grad_x, grad_y, grad_z) = ex.evaluate(back_func)(x_nd, y_nd, z_nd) tvm.testing.assert_allclose(forward.asnumpy(), expected_forward) tvm.testing.assert_allclose(grad_x.asnumpy(), np.ones_like(grad_x.asnumpy())) tvm.testing.assert_allclose(grad_y.asnumpy(), np.ones_like(grad_y.asnumpy())) tvm.testing.assert_allclose(grad_z.asnumpy(), -1 * np.ones_like(grad_z.asnumpy()))
def test_square_second_order(): shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) func = relay.Function([x], x * x) back_func = relay.ir_pass.infer_type(gradient(func)) y = relay.var("y", t) back_func_adjusted = relay.Function([y], relay.TupleGetItem(relay.TupleGetItem(back_func(y), 1), 0)) back_func_adjusted = relay.ir_pass.infer_type(back_func_adjusted) back_back_func = relay.ir_pass.infer_type(gradient(back_func_adjusted)) assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])])) x_nd = rand(dtype, *shape) ex = create_executor() forward, (grad_x,) = ex.evaluate(back_back_func)(x_nd) tvm.testing.assert_allclose(forward.asnumpy(), 2 * x_nd.asnumpy()) tvm.testing.assert_allclose(grad_x.asnumpy(), 2 * np.ones_like(grad_x.asnumpy()))
def test_empty_ad(): shape = (10, 10) dtype = "float32" t = TensorType(shape, dtype) d = Var("d", t) f = Function([d], d) g = dcpe(gradient(f)) expected = Function([d], Tuple([d, Tuple([op.ones_like(d)])])) assert alpha_equal(g, expected)
def test_ad(): shape = (10, 10) dtype = "float32" t = relay.TensorType(shape, dtype) d = relay.Var("d", t) f = relay.Function([d], d * d) g = dcpe(gradient(f)) m = d * d o = relay.op.ones_like(m) grad = relay.op.zeros_like(d) + relay.op.collapse_sum_like( o * d, d) + relay.op.collapse_sum_like(o * d, d) expected = relay.Function([d], relay.Tuple([m, relay.Tuple([grad])])) assert alpha_equal(g, expected)
def test_ad(): # TODO(MK): fix me shape = (10, 10) dtype = "float32" t = relay.TensorType(shape, dtype) d = relay.Var("d", t) f = relay.Function([d], d * d) g = dcpe(gradient(f)) m = d * d o = relay.op.ones_like(m) grad = relay.op.zeros_like(d) + relay.op.collapse_sum_like(o * d, d) + relay.op.collapse_sum_like(o * d, d) expected = relay.Function([d], relay.Tuple([m, relay.Tuple([grad])])) assert alpha_equal(g, expected)
def test_sub(): shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) func = relay.Function([x], x - x) back_func = relay.ir_pass.infer_type(gradient(func)) assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])])) ex = create_executor() x = rand(dtype, *shape) forward, (grad,) = ex.evaluate(back_func)(x) tvm.testing.assert_allclose(forward.asnumpy(), np.zeros_like(x.asnumpy())) tvm.testing.assert_allclose(grad.asnumpy(), np.zeros_like(x.asnumpy()))
def test_ad(): shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) func = relay.Function([x], x + x) back_func = relay.ir_pass.infer_type(gradient(func)) feats = detect_feature(back_func) assert feats == set([ Feature.fVar, Feature.fTuple, Feature.fTupleGetItem, Feature.fFunction, Feature.fOp, Feature.fCall, Feature.fLet, Feature.fRefCreate, Feature.fRefRead, Feature.fRefWrite ])
def dcpe(expr, mod=None, grad=False): passes = [ transform.PartialEvaluate(), transform.DeadCodeElimination(inline_once=True) ] if grad: expr = gradient(expr) if mod: assert isinstance(expr, Function) mod[mod.entry_func] = expr seq = transform.Sequential(passes) mod = seq(mod) return mod[mod.entry_func] return transform.OptimizeOnExpr(expr, passes)
def test_pow(): mod = relay.Module() p = Prelude(mod) 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([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) tvm.testing.assert_allclose(forward.asnumpy(), 8 * i_nd.asnumpy()) tvm.testing.assert_allclose(grad_i.asnumpy(), 8 * np.ones_like(grad_i.asnumpy()))
def check_single_op(opfunc, ref): shape = (10, 4) dtype = 'float32' tp = relay.TensorType(shape, dtype) x = relay.var("x", tp) y = opfunc(x) if ref is not None: data = np.random.rand(*shape).astype(dtype) ref_grad = ref(data) fwd_func = relay.Function([x], y) bwd_func = infer_type(gradient(fwd_func)) for target, ctx in ctx_list(): intrp = relay.create_executor(ctx=ctx, target=target) op_res, (op_grad, ) = intrp.evaluate(bwd_func)(data) np.testing.assert_allclose(op_grad.asnumpy(), ref_grad, rtol=0.01)
def test_ad(): shape = (10, 10) dtype = "float32" t = TensorType(shape, dtype) d = Var("d", t) f = Function([d], d * d) g = dcpe(gradient(f)) m = d * d x = relay.Var("x") o = op.ones_like(x) x1 = relay.Var("x1") grad = op.zeros_like(d) + op.collapse_sum_like( x1 * d, d) + op.collapse_sum_like(x1 * d, d) body = Tuple([x, Tuple([grad])]) body = relay.Let(x1, o, body) expected = Function([d], relay.Let(x, m, body)) assert alpha_equal(g, expected)
def check_binary_op(opfunc, ref): s = (5, 10, 5) t = relay.TensorType((5, 10, 5)) x = relay.var("x", t) y = relay.var("y", t) z = opfunc(x, y) x_data = np.random.rand(*s).astype(t.dtype) y_data = np.random.rand(*s).astype(t.dtype) ref_grad0, ref_grad1 = ref(x_data, y_data) fwd_func = relay.Function([x, y], z) bwd_func = infer_type(gradient(fwd_func)) for target, ctx in ctx_list(): intrp = relay.create_executor(ctx=ctx, target=target) op_res, (op_grad0, op_grad1) = intrp.evaluate(bwd_func)(x_data, y_data) np.testing.assert_allclose(op_grad0.asnumpy(), ref_grad0, rtol=0.01) np.testing.assert_allclose(op_grad1.asnumpy(), ref_grad1, rtol=0.01)
def test_ref(): shape = (10, 10) dtype = 'float32' t = relay.TensorType(shape, dtype) x = relay.var("x", t) r = relay.Var("r") u = relay.Var("u") body = relay.RefRead(r) body = relay.Let(u, relay.RefWrite(r, relay.RefRead(r) + relay.RefRead(r)), body) body = relay.Let(r, relay.RefCreate(x), body) func = relay.Function([x], body) back_func = relay.ir_pass.infer_type(gradient(func)) assert back_func.checked_type == relay.FuncType([t], relay.TupleType([t, relay.TupleType([t])])) x_nd = rand(dtype, *shape) ex = create_executor() forward, (grad_x,) = ex.evaluate(back_func)(x_nd) tvm.testing.assert_allclose(forward.asnumpy(), 2 * x_nd.asnumpy()) tvm.testing.assert_allclose(grad_x.asnumpy(), 2 * np.ones_like(grad_x.asnumpy()))