()
| 59 | |
| 60 | |
| 61 | def test_l1_loss_append(): |
| 62 | s = Module["forward"].ret_struct_info |
| 63 | l1_loss = relax.training.loss.L1Loss(reduction="sum") |
| 64 | After = relax.training.AppendLoss("forward", l1_loss(s, s), l1_loss.num_backbone_outputs)( |
| 65 | Module |
| 66 | ) |
| 67 | |
| 68 | @R.function |
| 69 | def expected( |
| 70 | x: R.Tensor((2, 4), "float32"), |
| 71 | w: R.Tensor((4, 4), "float32"), |
| 72 | b: R.Tensor((2, 4), "float32"), |
| 73 | targets: R.Tensor((2, 4), "float32"), |
| 74 | ) -> R.Tensor((), "float32"): |
| 75 | R.func_attr({"global_symbol": "forward_loss"}) |
| 76 | with R.dataflow(): |
| 77 | lv: R.Tensor((2, 4), "float32") = R.matmul(x, w, out_dtype="") |
| 78 | out: R.Tensor((2, 4), "float32") = R.add(lv, b) |
| 79 | lv1: R.Tensor((2, 4), "float32") = R.subtract(out, targets) |
| 80 | lv11: R.Tensor((2, 4), "float32") = R.abs(lv1) |
| 81 | gv: R.Tensor((), "float32") = R.sum(lv11, axis=None, keepdims=False) |
| 82 | R.output(gv) |
| 83 | return gv |
| 84 | |
| 85 | assert_structural_equal(After["forward_loss"], expected) |
| 86 | |
| 87 | |
| 88 | def test_mse_loss(): |
nothing calls this directly
no test coverage detected
searching dependent graphs…