Function test_emit

tests/python/relax/test_testing_nn.py:27–47

()

def test_emit():
    class ReLU(nn.Module):
        def forward(self, input: relax.Expr) -> relax.Var:
            return nn.emit(relax.op.nn.relu(input))

    @I.ir_module
    class Expected:
        @R.function
        def main(x: R.Tensor((32, 32), dtype="float32")) -> R.Tensor((32, 32), dtype="float32"):
            gv: R.Tensor((32, 32), dtype="float32") = R.nn.relu(x)
            return gv

    bb = relax.BlockBuilder()
    with bb.function("main"):
        model = ReLU()
        x = nn.Placeholder((32, 32), dtype="float32", name="x")
        output = model(x)
        params = [x] + model.parameters()
        bb.emit_func_output(output, params)

    tvm.ir.assert_structural_equal(bb.get(), Expected)


def test_get_param():