ckt624 commented on a change in pull request #15349: Numpy Tensordot Operator
URL: https://github.com/apache/incubator-mxnet/pull/15349#discussion_r302023553
##########
File path: python/mxnet/ndarray/numpy/_op.py
##########
@@ -28,7 +28,88 @@
__all__ = ['zeros', 'ones', 'maximum', 'minimum', 'stack', 'arange', 'argmax',
'add', 'subtract', 'multiply', 'divide', 'mod', 'power',
'concatenate',
'clip', 'split', 'swapaxes', 'expand_dims', 'tile', 'linspace',
- 'sin', 'cos', 'sinh', 'cosh', 'log10', 'sqrt']
+ 'sin', 'cos', 'sinh', 'cosh', 'log10', 'sqrt', 'tensordot']
+
+@set_module('mxnet.ndarray.numpy')
+def tensordot(a, b, axes = 2):
+ """
+ tensordot(a, b, axes=2)
+
+ Compute tensor dot product along specified axes for arrays >= 1-D.
+
+ Given two tensors (arrays of dimension greater than or equal to one),
+ `a` and `b`, and an array_like object containing two array_like
+ objects, ``(a_axes, b_axes)``, sum the products of `a`'s and `b`'s
+ elements (components) over the axes specified by ``a_axes`` and
+ ``b_axes``. The third argument can be a single non-negative
+ integer_like scalar, ``N``; if it is such, then the last ``N``
+ dimensions of `a` and the first ``N`` dimensions of `b` are summed
+ over.
+
+ Parameters
+ ----------
+ a, b : array_like, len(shape) >= 1
+ Tensors to "dot".
+
+ axes : int or (2,) array_like
+ * integer_like
+ If an int N, sum over the last N axes of `a` and the first N axes
+ of `b` in order. The sizes of the corresponding axes must match.
+ * (2,) array_like
+ Or, a list of axes to be summed over, first sequence applying to `a`,
+ second to `b`. Both elements array_like must be of the same length.
+
+ See Also
+ --------
+ dot, einsum
+
+ Notes
+ -----
+ Three common use cases are:
+ * ``axes = 0`` : tensor product :math:`a\otimes b`
+ * ``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 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.
+
+ Examples
+ --------
+ >>> 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]))
+ >>> c.shape
+ (5, 2)
+ >>> c
+ array([[ 4400., 4730.],
+ [ 4532., 4874.],
Review comment:
Fix. Thx.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
