Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

DOC: Fix a typo in description and add an example of numpy.tensordot #23547

Open
wants to merge 13 commits into
base: main
Choose a base branch
from
31 changes: 26 additions & 5 deletions numpy/core/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -971,22 +971,42 @@ def tensordot(a, b, axes=2):
* ``axes = 1`` : tensor dot product :math:`a\\cdot b`
* ``axes = 2`` : (default) tensor double contraction :math:`a:b`

When `axes` is integer_like, the sequence for evaluation will be: first
the -Nth axis in `a` and 0th axis in `b`, and the -1th axis in `a` and
Nth axis in `b` last.
When `axes` is integer_like, the sequence of axes for evaluation
will be: from the -Nth axis to the -1th axis in `a`,
and from the 0th axis to (N-1)th axis in `b`.
For example, ``axes = 2`` is the equal to
``axes = [[-2, -1], [0, 1]]``.
When N-1 is smaller than 0, or when -N is larger than -1,
the element of 'a' and 'b' are defined as the 'axes'.
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agree with @mattip, this part was a little unclear and redundant before.
Therefore, I add some details to help express more clearly.


When there is more than one axis to sum over - and they are not the last
(first) axes of `a` (`b`) - the argument `axes` should consist of
two sequences of the same length, with the first axis to sum over given
first in both sequences, the second axis second, and so forth.
The calculation can be referred to enisum().
John-Venti marked this conversation as resolved.
Show resolved Hide resolved

The shape of the result consists of the non-contracted axes of the
first tensor, followed by the non-contracted axes of the second.

Examples
--------
A "traditional" example:

"traditional" examples:
>>> # An example on integer_like
>>> a_0 = np.array([[1, 2], [3, 4]])
>>> b_0 = np.array([[5, 6], [7, 8]])
>>> c_0 = np.tensordot(a_0, b_0, axes=0)
>>> c_0.shape
(2, 2, 2, 2)
>>> c_0
array([[[[ 5, 6],
[ 7, 8]],
[[10, 12],
[14, 16]]],
[[[15, 18],
[21, 24]],
[[20, 24],
[28, 32]]]])
>>> # An example on array_like
John-Venti marked this conversation as resolved.
Show resolved Hide resolved
>>> a = np.arange(60.).reshape(3,4,5)
>>> b = np.arange(24.).reshape(4,3,2)
>>> c = np.tensordot(a,b, axes=([1,0],[0,1]))
Expand All @@ -998,6 +1018,7 @@ def tensordot(a, b, axes=2):
[4664., 5018.],
[4796., 5162.],
[4928., 5306.]])

>>> # A slower but equivalent way of computing the same...
John-Venti marked this conversation as resolved.
Show resolved Hide resolved
>>> d = np.zeros((5,2))
>>> for i in range(5):
Expand Down