import paddle

paddle.enable_static()

def fn_1():
    return paddle.full(shape=[1, 2], dtype='float32', fill_value=1)

def fn_2():
    return paddle.full(shape=[2, 2], dtype='int32', fill_value=2)

def fn_3():
    return paddle.full(shape=[3], dtype='int32', fill_value=3)

main_program = paddle.static.default_startup_program()
startup_program = paddle.static.default_main_program()

with paddle.static.program_guard(main_program, startup_program):
    x = paddle.full(shape=[1], dtype='float32', fill_value=0.3)
    y = paddle.full(shape=[1], dtype='float32', fill_value=0.1)
    z = paddle.full(shape=[1], dtype='float32', fill_value=0.2)

    pred_1 = paddle.less_than(z, x)  # true: 0.2 < 0.3
    pred_2 = paddle.less_than(x, y)  # false: 0.3 < 0.1
    pred_3 = paddle.equal(x, y)      # false: 0.3 == 0.1

    # Call fn_1 because pred_1 is True
    out_1 = paddle.static.nn.case(
        pred_fn_pairs=[(pred_1, fn_1), (pred_2, fn_2)], default=fn_3)

    # Argument default is None and no pred in pred_fn_pairs is True. fn_3 will be called.
    # because fn_3 is the last callable in pred_fn_pairs.
    out_2 = paddle.static.nn.case(pred_fn_pairs=[(pred_2, fn_2), (pred_3, fn_3)])

    exe = paddle.static.Executor(paddle.CPUPlace())
    res_1, res_2 = exe.run(main_program, fetch_list=[out_1, out_2])
    print(res_1)  # [[1. 1.]]
    print(res_2)  # [3 3 3]