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
36 changes: 30 additions & 6 deletions numpy/core/numeric.py
Original file line number Diff line number Diff line change
Expand Up @@ -971,21 +971,43 @@ 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`.

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 ``numpy.einsum``.

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:
--------
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

>>> a = np.arange(60.).reshape(3,4,5)
>>> b = np.arange(24.).reshape(4,3,2)
Expand All @@ -998,7 +1020,9 @@ def tensordot(a, b, axes=2):
[4664., 5018.],
[4796., 5162.],
[4928., 5306.]])
>>> # A slower but equivalent way of computing the same...

A slower but equivalent way of computing the same...

>>> d = np.zeros((5,2))
>>> for i in range(5):
... for j in range(2):
Expand Down