Source code for mars.tensor.linalg.tensordot

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2021 Alibaba Group Holding Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
from import Iterable

import numpy as np

from ... import opcodes as OperandDef
from ...serialization.serializables import FieldTypes, KeyField, TupleField
from ...utils import has_unknown_shape
from ..utils import unify_chunks
from ..array_utils import as_same_device, device, is_sparse_module
from ..operands import TensorOperand, TensorOperandMixin
from ..arithmetic.utils import chunk_tree_add
from ..datasource import tensor as astensor
from ..core import TensorOrder

class TensorTensorDot(TensorOperand, TensorOperandMixin):
    _op_type_ = OperandDef.TENSORDOT

    _a = KeyField("a")
    _b = KeyField("b")
    _a_axes = TupleField("a_axes", FieldTypes.int32)
    _b_axes = TupleField("b_axes", FieldTypes.int32)

    def __init__(self, a_axes=None, b_axes=None, **kw):
        super().__init__(_a_axes=a_axes, _b_axes=b_axes, **kw)

    def a(self):
        return self._a

    def b(self):
        return self._b

    def a_axes(self):
        return self._a_axes

    def b_axes(self):
        return self._b_axes

    def _set_inputs(self, inputs):
        self._a = self._inputs[0]
        self._b = self._inputs[1]

    def __call__(self, a, b):
        shape = tuple(
            s for i, s in enumerate(a.shape) if i not in set(self._a_axes)
        ) + tuple(s for i, s in enumerate(b.shape) if i not in set(self._b_axes))
        return self.new_tensor([a, b], shape, order=TensorOrder.C_ORDER)

    def estimate_size(cls, ctx, op):
        chunk = op.outputs[0]
        if chunk.is_sparse():
            return super().estimate_size(ctx, op)

        # empirical value in real environments
        calc_usage = chunk.nbytes

        # add input sizes when sparse-to-dense is needed
        for inp in chunk.inputs:
            if inp.is_sparse():
                calc_usage += inp.nbytes

        ctx[chunk.key] = (chunk.nbytes, calc_usage)

    def tile(cls, op):
        a, b, a_axes, b_axes = op.a, op.b, op.a_axes, op.b_axes

        c = itertools.count(max(a.ndim, b.ndim))
        a_ax = tuple(a_axes.index(i) if i in a_axes else next(c) for i in range(a.ndim))
        b_ax = tuple(b_axes.index(i) if i in b_axes else next(c) for i in range(b.ndim))
        if has_unknown_shape(*op.inputs):
        a, b = yield from unify_chunks((a, a_ax), (b, b_ax))
        out = op.outputs[0]

        a_output_indexes = [
            range(len(a.nsplits[i])) for i in range(a.ndim) if i not in a_axes
        b_output_indexes = [
            range(len(b.nsplits[i])) for i in range(b.ndim) if i not in b_axes
        output_axes = [(0, i) for i in range(a.ndim) if i not in a_axes] + [
            (1, i) for i in range(b.ndim) if i not in b_axes

        out_chunks = []
        for out_idx in itertools.product(
            *itertools.chain(a_output_indexes, b_output_indexes)
            a_indexes = [None] * a.ndim
            b_indexes = [None] * b.ndim
            tensor_shape = []
            for i, idx in enumerate(out_idx):
                t_idx, axis = output_axes[i]
                t = (a, b)[t_idx]
                (a_indexes if t_idx == 0 else b_indexes)[axis] = idx
            tensor_shape = tuple(tensor_shape)

            tensordot_chunks = []
            for contract_indexes in itertools.product(
                *[range(len(a.nsplits[ax])) for ax in a_axes]
                a_indices, b_indices = list(a_indexes), list(b_indexes)
                for a_axis, contract_index in zip(a_axes, contract_indexes):
                    a_indices[a_axis] = contract_index
                for b_axis, contract_index in zip(b_axes, contract_indexes):
                    b_indices[b_axis] = contract_index

                tensordot_chunk_op = op.copy().reset_key()
                tensordot_chunk = tensordot_chunk_op.new_chunk(
                    [a.cix[tuple(a_indices)], b.cix[tuple(b_indices)]],

            if len(tensordot_chunks) == 1:
                c = tensordot_chunks[0]
                chunk_op = c.op.copy()
                chunk = chunk_op.new_chunk(
                    c.inputs, shape=c.shape, index=out_idx, order=out.order
                chunk = chunk_tree_add(
                    op.dtype, tensordot_chunks, out_idx, tensor_shape, sparse=op.sparse

        get_nsplits = lambda t_idx, i: (a, b)[t_idx].nsplits[i]
        nsplits = [get_nsplits(*it) for it in output_axes]
        new_op = op.copy()
        return new_op.new_tensors([a, b], out.shape, chunks=out_chunks, nsplits=nsplits)

    def execute(cls, ctx, op):
        (a, b), device_id, xp = as_same_device(
            [ctx[c.key] for c in op.inputs], device=op.device, ret_extra=True

        axes = op.a_axes, op.b_axes
        with device(device_id):
            if not op.sparse and is_sparse_module(xp):
                # tell sparse to do calculation on numpy or cupy dot
                ctx[op.outputs[0].key] = xp.tensordot(a, b, axes, sparse=False)
                ret = xp.tensordot(a, b, axes)
                out = op.outputs[0]
                ctx[out.key] = ret.astype(ret.dtype, order=out.order.value, copy=False)

[docs]def tensordot(a, b, axes=2, sparse=None): """ Compute tensor dot product along specified axes for tensors >= 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 -------- >>> import mars.tensor as mt A "traditional" example: >>> a = mt.arange(60.).reshape(3,4,5) >>> b = mt.arange(24.).reshape(4,3,2) >>> c = mt.tensordot(a,b, axes=([1,0],[0,1])) >>> c.shape (5, 2) >>> r = c.execute() >>> r array([[ 4400., 4730.], [ 4532., 4874.], [ 4664., 5018.], [ 4796., 5162.], [ 4928., 5306.]]) >>> # A slower but equivalent way of computing the same... >>> ra = np.arange(60.).reshape(3,4,5) >>> rb = np.arange(24.).reshape(4,3,2) >>> d = np.zeros((5,2)) >>> for i in range(5): ... for j in range(2): ... for k in range(3): ... for n in range(4): ... d[i,j] += ra[k,n,i] * rb[n,k,j] >>> r == d array([[ True, True], [ True, True], [ True, True], [ True, True], [ True, True]], dtype=bool) An extended example taking advantage of the overloading of + and \\*: >>> a = mt.array(range(1, 9)) >>> a.shape = (2, 2, 2) >>> A = mt.array(('a', 'b', 'c', 'd'), dtype=object) >>> A.shape = (2, 2) >>> a.execute(); A.execute() array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) array([[a, b], [c, d]], dtype=object) >>> mt.tensordot(a, A).execute() # third argument default is 2 for double-contraction array([abbcccdddd, aaaaabbbbbbcccccccdddddddd], dtype=object) >>> mt.tensordot(a, A, 1).execute() array([[[acc, bdd], [aaacccc, bbbdddd]], [[aaaaacccccc, bbbbbdddddd], [aaaaaaacccccccc, bbbbbbbdddddddd]]], dtype=object) >>> mt.tensordot(a, A, 0).execute() # tensor product (result too long to incl.) array([[[[[a, b], [c, d]], ... >>> mt.tensordot(a, A, (0, 1)).execute() array([[[abbbbb, cddddd], [aabbbbbb, ccdddddd]], [[aaabbbbbbb, cccddddddd], [aaaabbbbbbbb, ccccdddddddd]]], dtype=object) >>> mt.tensordot(a, A, (2, 1)).execute() array([[[abb, cdd], [aaabbbb, cccdddd]], [[aaaaabbbbbb, cccccdddddd], [aaaaaaabbbbbbbb, cccccccdddddddd]]], dtype=object) >>> mt.tensordot(a, A, ((0, 1), (0, 1))).execute() array([abbbcccccddddddd, aabbbbccccccdddddddd], dtype=object) >>> mt.tensordot(a, A, ((2, 1), (1, 0))).execute() array([acccbbdddd, aaaaacccccccbbbbbbdddddddd], dtype=object) """ a = astensor(a) b = astensor(b) if isinstance(axes, Iterable): a_axes, b_axes = axes else: a_axes = tuple(range(a.ndim - 1, a.ndim - axes - 1, -1)) b_axes = tuple(range(0, axes)) if isinstance(a_axes, Iterable): a_axes = tuple(a_axes) else: a_axes = (a_axes,) a_axes = tuple(axis if axis >= 0 else a.ndim + axis for axis in a_axes) if isinstance(b_axes, Iterable): b_axes = tuple(b_axes) else: b_axes = (b_axes,) b_axes = tuple(axis if axis >= 0 else b.ndim + axis for axis in b_axes) if ( a.shape and b.shape and not np.array_equal( np.array(a.shape)[list(a_axes)], np.array(b.shape)[list(b_axes)] ) ): raise ValueError("shape-mismatch for sum") sparse = sparse if sparse is not None else a.issparse() and b.issparse() op = TensorTensorDot( a_axes=a_axes, b_axes=b_axes, dtype=np.promote_types(a.dtype, b.dtype), sparse=sparse, ) return op(a, b)