def METHOD_NAME(distr, slice_axis_args, expected_axis_length):
    axis, begin, end = slice_axis_args
    distr_sliced = distr.slice_axis(axis, begin, end)
    assert distr_sliced.batch_shape[axis] == expected_axis_length