(target, dev)
| 140 | |
| 141 | @tvm.testing.parametrize_targets("llvm") |
| 142 | def test_complex(target, dev): |
| 143 | cst = relax.const(np.ones((6,)), dtype="float32") |
| 144 | cst1 = relax.const(np.array(3), dtype="int64") |
| 145 | |
| 146 | @tvm.script.ir_module |
| 147 | class Before: |
| 148 | @R.function |
| 149 | def main(x: R.Tensor((6,), "float32"), y: R.Tensor((6, 3, 4), "float32")): |
| 150 | with R.dataflow(): |
| 151 | lv1 = R.split(x, 2) |
| 152 | lv2 = lv1[0] |
| 153 | lv3 = lv1[1] |
| 154 | lv4 = lv2 + lv3 |
| 155 | lv5 = (lv4, lv3) |
| 156 | lv6 = R.concat(lv5) |
| 157 | lv7 = (x, x) |
| 158 | lv8 = R.concat(lv7) |
| 159 | lv9 = R.concat(lv7) |
| 160 | lv10 = R.add(lv8, lv9) |
| 161 | lv11 = R.split(lv10, 2) |
| 162 | lv12 = R.add(lv6, lv11[0]) |
| 163 | lv13 = cst |
| 164 | lv14 = R.add(lv12, lv13) |
| 165 | lv15 = R.subtract(lv13, lv14) |
| 166 | lv16 = R.multiply(lv14, lv15) |
| 167 | lv17 = R.multiply(lv15, lv16) |
| 168 | lv18 = R.tanh(lv17) |
| 169 | lv19 = R.sigmoid(lv18) |
| 170 | lv20 = R.permute_dims(y, axes=[0, 2, 1]) |
| 171 | lv21 = R.sigmoid(lv20) |
| 172 | lv22 = R.matmul(y, lv21) |
| 173 | lv23 = R.sum(lv22, axis=[1, 2]) |
| 174 | lv24 = R.add(lv19, lv23) |
| 175 | lv25 = R.nn.log_softmax(lv24) |
| 176 | gv = R.nn.nll_loss(lv25, cst1) |
| 177 | R.output(gv) |
| 178 | return gv |
| 179 | |
| 180 | After = relax.transform.Gradient("main")(Before) |
| 181 | args = [] |
| 182 | for arg in After["main_adjoint"].params: |
| 183 | shape = [int(l) for l in arg.struct_info.shape] |
| 184 | args.append(rand("float32", *shape)) |
| 185 | |
| 186 | vm_before = _legalize_and_build(Before, target, dev) |
| 187 | vm_after = _legalize_and_build(After, target, dev) |
| 188 | _, grad = vm_after["main_adjoint"](*args) |
| 189 | |
| 190 | def func(*inputs): |
| 191 | loss = vm_before["main"](*[tvm.runtime.tensor(i) for i in inputs]) |
| 192 | return loss.numpy() |
| 193 | |
| 194 | check_numerical_grads(func, [i.numpy() for i in args], [i.numpy() for i in grad]) |
| 195 | |
| 196 | |
| 197 | @tvm.testing.parametrize_targets("llvm") |
nothing calls this directly
no test coverage detected
searching dependent graphs…