diff options
Diffstat (limited to 'numpy/core/einsumfunc.py')
| -rw-r--r-- | numpy/core/einsumfunc.py | 9 |
1 files changed, 3 insertions, 6 deletions
diff --git a/numpy/core/einsumfunc.py b/numpy/core/einsumfunc.py index d9f88cb1c..963c696ae 100644 --- a/numpy/core/einsumfunc.py +++ b/numpy/core/einsumfunc.py @@ -888,9 +888,8 @@ def einsum_path(*operands, **kwargs): broadcast_indices = [set(x) for x in broadcast_indices] # Compute size of each input array plus the output array - size_list = [] - for term in input_list + [output_subscript]: - size_list.append(_compute_size_by_dict(term, dimension_dict)) + size_list = [_compute_size_by_dict(term, dimension_dict) + for term in input_list + [output_subscript]] max_size = max(size_list) if memory_limit is None: @@ -1375,9 +1374,7 @@ def einsum(*operands, **kwargs): # Start contraction loop for num, contraction in enumerate(contraction_list): inds, idx_rm, einsum_str, remaining, blas = contraction - tmp_operands = [] - for x in inds: - tmp_operands.append(operands.pop(x)) + tmp_operands = [operands.pop(x) for x in inds] # Do we need to deal with the output? handle_out = specified_out and ((num + 1) == len(contraction_list)) |