Function test_reduce_axis_names

dali/test/python/operator_2/test_reduce.py:476–525

()

def test_reduce_axis_names():
    reductions = [
        fn.reductions.max,
        fn.reductions.min,
        fn.reductions.mean,
        fn.reductions.mean_square,
        fn.reductions.sum,
        fn.reductions.rms,
    ]

    reductions_with_mean_input = [fn.reductions.std_dev, fn.reductions.variance]

    batch_fn = Batch3D(np.float32)
    batch_size = batch_fn.batch_size()

    def get_batch():
        return batch_fn()

    axes_and_names = [
        ((), ""),
        (0, "A"),
        (1, "B"),
        (2, "C"),
        ((0, 1), "AB"),
        ((0, 2), "AC"),
        ((1, 2), "BC"),
        ((0, 1, 2), "ABC"),
    ]

    for axes, axis_names in axes_and_names:
        for reduction in reductions:
            yield (
                run_reduce_with_layout,
                batch_size,
                get_batch,
                reduction,
                axes,
                axis_names,
                batch_fn,
            )
        for reduction in reductions_with_mean_input:
            yield (
                run_reduce_with_layout_with_mean_input,
                batch_size,
                get_batch,
                reduction,
                axes,
                axis_names,
                batch_fn,
            )


_random_buf = None