Function test_empty_dfb_after_removal

tests/python/relax/test_binding_rewrite.py:193–206

()

def test_empty_dfb_after_removal():
    root_fn = DeadDFBlock["main"]
    dfb = root_fn.body.blocks[0]

    rwt = DataflowBlockRewrite(dfb, root_fn)
    rwt.remove_unused(DeadDFBlock["main"].body.blocks[0].bindings[0].var)

    @tvm.script.ir_module
    class GroundTruth:
        @R.function
        def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
            return x

    tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])


def test_empty_dfb_after_all_removal():