Function conditional_wrapper

dali/test/python/test_dali_variable_batch_size.py:1592–1617

(max_batch_size, input_data, device)

def test_conditional():
    def conditional_wrapper(max_batch_size, input_data, device):
        @experimental_pipeline_def(
            enable_conditionals=True,
            batch_size=max_batch_size,
            num_threads=4,
            device_id=0,
        )
        def actual_pipe():
            variable_condition = fn.external_source(source=input_data, cycle=False, device=device)
            variable_data = variable_condition + 42.0
            if variable_condition:
                other_variable_data = variable_condition + 100
                output = variable_data + other_variable_data
            else:
                output = types.Constant(np.array(42.0), device="cpu")
            logical_expr = variable_condition or not variable_condition
            logical_expr2 = not variable_condition and variable_condition
            return (
                output,
                variable_condition,
                variable_data,
                logical_expr,
                logical_expr2,
            )

        return actual_pipe()

    check_pipeline(
        generate_data(31, 13, custom_shape_generator(), lo=False, hi=True, dtype=np.bool_),