Function _test_seq_input

dali/test/python/sequences_test_utils.py:267–314

(num_iters, operator_fn, fixed_params, input_params, input_data: ArgData, rng)

def _test_seq_input(num_iters, operator_fn, fixed_params, input_params, input_data: ArgData, rng):
    @pipeline_def
    def pipeline(args_data: List[ArgData]):
        pos_args = [arg_data for arg_data in args_data if arg_data.desc.is_positional_arg]
        pos_nodes = [None] * len(pos_args)
        for arg_data in pos_args:
            assert 0 <= arg_data.desc.name < len(pos_nodes)
            assert pos_nodes[arg_data.desc.name] is None
            pos_nodes[arg_data.desc.name] = arg_data_node(arg_data)
        named_args = [arg_data for arg_data in args_data if not arg_data.desc.is_positional_arg]
        arg_nodes = {arg_data.desc.name: arg_data_node(arg_data) for arg_data in named_args}
        output = operator_fn(*pos_nodes, **fixed_params, **arg_nodes)
        return output

    assert num_iters >= len(input_data.data)
    max_batch_size = max(len(batch) for batch in input_data.data)

    params_provider = (
        input_params
        if isinstance(input_params, ParamsProviderBase)
        else ParamsProvider(input_params)
    )
    params_provider.setup(input_data, fixed_params, rng)
    args_data = params_provider.compute_params()
    seq_pipe = pipeline(
        args_data=[input_data, *args_data], batch_size=max_batch_size, num_threads=4, device_id=0
    )
    unfolded_input = params_provider.unfold_input()
    expanded_args_data = params_provider.expand_params()
    max_uf_batch_size = max(len(batch) for batch in unfolded_input.data)
    baseline_pipe = pipeline(
        args_data=[unfolded_input, *expanded_args_data],
        batch_size=max_uf_batch_size,
        num_threads=4,
        device_id=0,
    )

    for _ in range(num_iters):
        (seq_batch,) = seq_pipe.run()
        (baseline_batch,) = baseline_pipe.run()
        unfolded_layout = params_provider.unfold_output_layout(seq_batch.layout())
        assert (
            unfolded_layout == baseline_batch.layout()
        ), f"'{unfolded_layout}' == '{baseline_batch.layout()}'"
        batch = params_provider.unfold_output(as_batch(seq_batch))
        baseline_batch = as_batch(baseline_batch)
        assert len(batch) == len(baseline_batch)
        check_batch(batch, baseline_batch, len(batch))


def get_input_arg_per_sample(input_data, param_cb, rng):