(target, dev)
| 38 | |
| 39 | @tvm.testing.parametrize_targets("llvm") |
| 40 | def test_manual_gradient(target, dev): |
| 41 | # The expression computed is sum((2x - 2y) * (y + z)) |
| 42 | # the gradient of x is broadcast_to(2y + 2z, x.shape) |
| 43 | # the gradient of y is collapse_sum_to((2x - 4y - 2z), y.shape) |
| 44 | # the gradient of z is collapse_sum_to((2x - 2y), z.shape) |
| 45 | # the gradient of u is 0 |
| 46 | @I.ir_module |
| 47 | class Before: |
| 48 | @R.function |
| 49 | def main( |
| 50 | x: R.Tensor((3, 5), "float32"), |
| 51 | y: R.Tensor((5,), "float32"), |
| 52 | z: R.Tensor((5,), "float32"), |
| 53 | u: R.Tensor((5,), "float32"), |
| 54 | ): |
| 55 | with R.dataflow(): |
| 56 | lv1 = R.add(x, x) |
| 57 | lv2 = R.subtract(lv1, y) |
| 58 | lv3 = R.subtract(lv2, y) |
| 59 | lv4 = R.add(y, z) |
| 60 | lv5 = R.multiply(lv3, lv4) |
| 61 | lv6 = R.sum(lv5) |
| 62 | R.output(lv6) |
| 63 | return lv6 |
| 64 | |
| 65 | After = relax.transform.Gradient("main")(Before) |
| 66 | |
| 67 | args = [rand("float32", 3, 5), rand("float32", 5), rand("float32", 5), rand("float32", 5)] |
| 68 | args_np = [x.numpy() for x in args] |
| 69 | |
| 70 | vm = _legalize_and_build(After, target, dev) |
| 71 | output, grads = vm["main_adjoint"](*args) |
| 72 | output_np = np.sum((2 * args_np[0] - 2 * args_np[1]) * (args_np[1] + args_np[2])) |
| 73 | assert_allclose(output.numpy(), output_np, atol=1e-4) |
| 74 | |
| 75 | expected_grads_nd = [ |
| 76 | (2 * args_np[1] + 2 * args_np[2]) * np.ones_like(args_np[0]), |
| 77 | np.sum((2 * args_np[0] - 4 * args_np[1] - 2 * args_np[2]), axis=0), |
| 78 | np.sum((2 * args_np[0] - 2 * args_np[1]), axis=0), |
| 79 | np.zeros_like(args_np[3]), |
| 80 | ] |
| 81 | for i, j in zip(grads, expected_grads_nd): |
| 82 | assert_allclose(i.numpy(), j, atol=1e-4) |
| 83 | |
| 84 | |
| 85 | @tvm.testing.parametrize_targets("llvm") |
nothing calls this directly
no test coverage detected
searching dependent graphs…