ENH: update numpy.linalg.multi_dot to accept an out argument
#15715
Conversation
|
I'm finding the PyPy pipeline test failures a bit cryptic Is there something special that needs to be done for PyPy? |
|
Not sure what is going on with this failure in the PyPy pipeline, although maybe we can rewrite that script. It fails with: |
|
For some reason there is pip3 in pypy/bin but not pip. Perhaps we should be building against a known good version of pypy and not latest HEAD. |
charris
added
01 - Enhancement
56 - Needs Release Note.
|
Needs a release note. |
Added one in ef7486a. Is it in the proper format? |
|
@eric-wieser @wkschwartz do you have any more thoughts on this? |
numpy/linalg/linalg.py
Outdated
| @@ -2610,12 +2610,12 @@ def norm(x, ord=None, axis=None, keepdims=False): | |||
|
|
|||
| # multi_dot | |||
|
|
|||
| def _multidot_dispatcher(arrays): | |||
| def _multidot_dispatcher(arrays, out=None): | |||
| return arrays | |||
There was a problem hiding this comment.
@eric-wieser @shoyer grepping around (rg -A 4 "_dispatcher.*out") it seems that out should be returned by the dispatcher. Other functions with an optional out return it. Added in 5427885
numpy/linalg/linalg.py
Outdated
| return arrays | ||
|
|
||
|
|
||
| @array_function_dispatch(_multidot_dispatcher) | ||
| def multi_dot(arrays): | ||
| def multi_dot(arrays, out=None): |
There was a problem hiding this comment.
Thoughts on making this as follows?
| def multi_dot(arrays, out=None): | |
| def multi_dot(arrays, *, out=None): |
There was a problem hiding this comment.
+1 let's make this keyword only!
There was a problem hiding this comment.
+1 from me as well.
normally I'd be concerned about consistency between functions, but for out it seems like such an improvement to have keyword-only that we should just do it.
There was a problem hiding this comment.
@eric-wieser @shoyer @rgommers Good point. Added in 1f26974
| `numpy.linalg.multi_dot` now accepts an ``out`` argument | ||
| -------------------------------------------------------- | ||
|
|
||
| ``out`` can be used to avoid creating unnecessary copies of the final product computed by `numpy.linalg.multidot`. |
| def test_basic_function_with_dynamic_programing_optimization_with_out_argument(self): | ||
| # multi_dot with four or more arguments uses the dynamic programing | ||
| # optimization and therefore deserve a separate | ||
| A = np.random.random((6, 2)) |
There was a problem hiding this comment.
comment seems cut off: "separate test"
These test name are quite long. Rather than test_basic_function_with can we have test_multidot_, and drop the _argument, so this would become test_multidot_with_dynamic_programming_and_out.
There was a problem hiding this comment.
I will commit your suggestion, but also removed the multidot_with part, this is already in the TestMultidot class, so that should be implicit.
I will then merge this later today or tomorrow.
mattip
removed
the
56 - Needs Release Note.
Co-Authored-By: Matti Picus <matti.picus@gmail.com>
|
What do we use for the concatenate dispatcher? That one could maybe be reused here. |
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>
|
Thanks @sslivkoff and reviewers |
outis a common argument in numpy functions that specifies where the result of the function should be stored. A typical use case foroutis performance-critical situations where it is desirable to avoid creating unnecessary copies of the result array.numpy.linalg.multi_dotis also commonly used in such situations, so it would be nice if it could accept anoutargument. This PR addsoutfunctionality tonumpy.linalg.multi_dot