@@ -1023,28 +1023,48 @@ def tensordot(a, b, axes=2):
1023
1023
Notes
1024
1024
-----
1025
1025
Three common use cases are:
1026
-
1027
- * ``axes = 0`` : tensor product :math:`a\\ otimes b`
1028
- * ``axes = 1`` : tensor dot product :math:`a\\ cdot b`
1029
- * ``axes = 2`` : (default) tensor double contraction :math:`a:b`
1030
-
1031
- When `axes` is a positive integer ``N``, the operation starts with
1032
- axis ``-N`` of `a` and axis ``0`` of `b`, and it continues through
1033
- axis ``-1`` of `a` and axis ``N-1`` of `b` (inclusive).
1026
+ * ``axes = 0`` : tensor product :math:`a\\ otimes b`
1027
+ * ``axes = 1`` : tensor dot product :math:`a\\ cdot b`
1028
+ * ``axes = 2`` : (default) tensor double contraction :math:`a:b`
1029
+
1030
+ When `axes` is integer_like, the sequence of axes for evaluation
1031
+ will be: from the -Nth axis to the -1th axis in `a`,
1032
+ and from the 0th axis to (N-1)th axis in `b`.
1033
+ For example, ``axes = 2`` is the equal to
1034
+ ``axes = [[-2, -1], [0, 1]]``.
1035
+ When N-1 is smaller than 0, or when -N is larger than -1,
1036
+ the element of `a` and `b` are defined as the `axes`.
1034
1037
1035
1038
When there is more than one axis to sum over - and they are not the last
1036
1039
(first) axes of `a` (`b`) - the argument `axes` should consist of
1037
1040
two sequences of the same length, with the first axis to sum over given
1038
1041
first in both sequences, the second axis second, and so forth.
1042
+ The calculation can be referred to ``numpy.einsum``.
1039
1043
1040
1044
The shape of the result consists of the non-contracted axes of the
1041
1045
first tensor, followed by the non-contracted axes of the second.
1042
1046
1043
1047
Examples
1044
- --------
1045
- A "traditional" example:
1048
+ --------
1049
+ An example on integer_like:
1050
+
1051
+ >>> a_0 = np.array([[1, 2], [3, 4]])
1052
+ >>> b_0 = np.array([[5, 6], [7, 8]])
1053
+ >>> c_0 = np.tensordot(a_0, b_0, axes=0)
1054
+ >>> c_0.shape
1055
+ (2, 2, 2, 2)
1056
+ >>> c_0
1057
+ array([[[[ 5, 6],
1058
+ [ 7, 8]],
1059
+ [[10, 12],
1060
+ [14, 16]]],
1061
+ [[[15, 18],
1062
+ [21, 24]],
1063
+ [[20, 24],
1064
+ [28, 32]]]])
1065
+
1066
+ An example on array_like:
1046
1067
1047
- >>> import numpy as np
1048
1068
>>> a = np.arange(60.).reshape(3,4,5)
1049
1069
>>> b = np.arange(24.).reshape(4,3,2)
1050
1070
>>> c = np.tensordot(a,b, axes=([1,0],[0,1]))
@@ -1056,7 +1076,9 @@ def tensordot(a, b, axes=2):
1056
1076
[4664., 5018.],
1057
1077
[4796., 5162.],
1058
1078
[4928., 5306.]])
1059
- >>> # A slower but equivalent way of computing the same...
1079
+
1080
+ A slower but equivalent way of computing the same...
1081
+
1060
1082
>>> d = np.zeros((5,2))
1061
1083
>>> for i in range(5):
1062
1084
... for j in range(2):
0 commit comments