def METHOD_NAME():
    dag, s = get_tiled_matmul()
    _, bufs = dag.apply_steps_from_state(s)
    assert bufs[1].shape[0] == 512
    assert bufs[1].shape[1] == 512
    _, bufs = dag.apply_steps_from_state(
        s, layout_rewrite=auto_scheduler.LayoutRewriteOption.REWRITE_FOR_PRE_TRANSFORMED
    )
    assert bufs[1].shape[0] == 4
    assert bufs[1].shape[1] == 8
    assert bufs[1].shape[2] == 4
    assert bufs[1].shape[3] == 4
    assert bufs[1].shape[4] == 512
    _, bufs = dag.apply_steps_from_state(
        s, layout_rewrite=auto_scheduler.LayoutRewriteOption.INSERT_TRANSFORM_STAGE
    )
    assert bufs[1].shape[0] == 512
    assert bufs[1].shape[1] == 512