[dynamo] fix functools.reduce() function with None as initial (#116398)

The `initial` argument in `functools.reduce` can be `None`.

```python
initial_missing = object()

def reduce(function, iterable, initial=initial_missing, /):
    it = iter(iterable)
    if initial is initial_missing:
        value = next(it)
    else:
        value = initial
    for element in it:
        value = function(value, element)
    return value
```

Reference:

- python/cpython#102759

Pull Request resolved: https://github.com/pytorch/pytorch/pull/116398
Approved by: https://github.com/Skylion007
This commit is contained in:
Xuehai Pan
2023-12-25 21:23:28 +00:00
committed by PyTorch MergeBot
parent c7e9c15102
commit 039fbeb016
3 changed files with 40 additions and 7 deletions

View File

@ -63,11 +63,19 @@ def func_with_default(a, b, some_default_arg=True):
return a - b
def make_test(fn):
def make_test(fn=None, expected_frame_count=1):
if fn is None:
return lambda fn: make_test(fn, expected_frame_count=expected_frame_count)
nargs = len(inspect.signature(fn).parameters)
def test_fn(self):
return torch._dynamo.testing.standard_test(self, fn=fn, nargs=nargs)
return torch._dynamo.testing.standard_test(
self,
fn=fn,
nargs=nargs,
expected_frame_count=expected_frame_count,
)
return test_fn
@ -870,6 +878,22 @@ class FunctionTests(torch._dynamo.test_case.TestCase):
def test_reduce(a, b, c, d):
return functools.reduce(operator.add, [a, b, c, d])
@make_test
def test_reduce_with_initial(a, b, c, d):
return functools.reduce(operator.add, [b, c, d], a)
@make_test(expected_frame_count=0)
def test_reduce_with_single(x):
return functools.reduce(lambda a, b: (a, b), [x])
@make_test(expected_frame_count=0)
def test_reduce_with_single_with_initial(x, y):
return functools.reduce(lambda a, b: (a, b), [y], x)
@make_test(expected_frame_count=0)
def test_reduce_with_none_initial(x):
return functools.reduce(lambda a, b: (a, b), [x], None)
@make_test
def test_tuple_contains(a, b):
v1 = "a"