Function arg_for_sample

dali/test/python/sequences_test_utils.py:331–340

(sample_idx, batch_idx, sample)

    assert frame_idx >= 0

    def arg_for_sample(sample_idx, batch_idx, sample):
        if check_broadcasting and rng.randint(1, 4) == 1:
            return np.array([param_cb(SampleDesc(rng, 0, sample_idx, batch_idx, sample))])
        num_frames = sample.shape[frame_idx]
        return np.array(
            [
                param_cb(SampleDesc(rng, frame_idx, sample_idx, batch_idx, sample))
                for frame_idx in range(num_frames)
            ]
        )

    return [
        [arg_for_sample(sample_idx, batch_idx, sample) for sample_idx, sample in enumerate(batch)]