diff --git a/tests/python/relay/test_pass_lambda_lift.py b/tests/python/relay/test_pass_lambda_lift.py new file mode 100644 index 000000000000..ffcdb5e3ea9c --- /dev/null +++ b/tests/python/relay/test_pass_lambda_lift.py @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import numpy as np +import pytest + +import tvm +from tvm import relay +from tvm.relay import transform + +def test_basic(): + mod = relay.Module() + x2 = relay.var('x2', shape=(10, 5)) + y2 = relay.var('y2', shape=(1, 5)) + level2_func = relay.Function([x2, y2], relay.op.add(x2, y2)) + + x1 = relay.var('x1', shape=(10, 5)) + y1 = relay.var('y1', shape=(1, 5)) + level1_func = relay.Function([x1, y1], level2_func(x1, y1)) + + mod["main"] = level1_func + new_mod = transform.LambdaLift()(mod) + assert len(new_mod.functions) == 2 + +if __name__ == "__main__": + pytest.main() +