Source code for mars.tensor.base.broadcast_to

# 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 numpy as np

from ... import opcodes as OperandDef
from ...serialization.serializables import KeyField, TupleField
from ..operands import TensorHasInput, TensorOperandMixin
from ..datasource import tensor as astensor
from ..array_utils import get_array_module, device

class TensorBroadcastTo(TensorHasInput, TensorOperandMixin):
    _op_type_ = OperandDef.BROADCAST_TO

    _input = KeyField("input")
    _shape = TupleField("shape")

    def __init__(self, shape=None, **kw):
        super().__init__(_shape=shape, **kw)

    def shape(self):
        return self._shape

    def __call__(self, tensor, shape):
        return self.new_tensor([tensor], shape)

    def tile(cls, op):
        tensor = op.outputs[0]
        in_tensor = op.inputs[0]
        shape = op.shape
        new_dim = tensor.ndim - in_tensor.ndim

        out_chunks = []
        for c in in_tensor.chunks:
            chunk_shape = shape[:new_dim] + tuple(
                s if in_tensor.shape[idx] != 1 else shape[new_dim + idx]
                for idx, s in enumerate(c.shape)
            chunk_idx = (0,) * new_dim + c.index
            chunk_op = op.copy().reset_key()
            chunk_op._shape = chunk_shape
            out_chunk = chunk_op.new_chunk(
                [c], shape=chunk_shape, index=chunk_idx, order=tensor.order

        nsplits = [
                for c in out_chunks
                if all(idx == 0 for j, idx in enumerate(c.index) if j != i)
            for i in range(len(out_chunks[0].shape))
        new_op = op.copy()
        return new_op.new_tensors(

    def execute(cls, ctx, op):
        xp = get_array_module(ctx[op.input.key])
        input_data = ctx[op.input.key]
        device_id = if hasattr(input_data, "device") else -1

        with device(device_id):
            shape = op.shape
            if any(np.isnan(s) for s in shape):
                shape = list(shape)
                new_dim = len(shape) - input_data.ndim
                for i in range(input_data.ndim):
                    if np.isnan(shape[i + new_dim]):
                        shape[i + new_dim] = input_data.shape[i]
            ctx[op.outputs[0].key] = xp.broadcast_to(input_data, shape)

[docs]def broadcast_to(tensor, shape): """Broadcast an tensor to a new shape. Parameters ---------- tensor : array_like The tensor to broadcast. shape : tuple The shape of the desired array. Returns ------- broadcast : Tensor Raises ------ ValueError If the tensor is not compatible with the new shape according to Mars's broadcasting rules. Examples -------- >>> import mars.tensor as mt >>> x = mt.array([1, 2, 3]) >>> mt.broadcast_to(x, (3, 3)).execute() array([[1, 2, 3], [1, 2, 3], [1, 2, 3]]) """ from ..core import Tensor tensor = tensor if isinstance(tensor, Tensor) else astensor(tensor) shape = tuple(shape) if isinstance(shape, (list, tuple)) else (shape,) if any(np.isnan(s) for s in tensor.shape): raise ValueError( "input tensor has unknown shape, need to call `.execute()` first" ) if tensor.shape == shape: return tensor new_ndim = len(shape) - tensor.ndim if new_ndim < 0: raise ValueError( "input operand has more dimensions than allowed by the axis remapping" ) if any(o != n for o, n in zip(tensor.shape, shape[new_ndim:]) if o != 1): raise ValueError( "operands could not be broadcast together " f"with remapped shapes [original->remapped]: {tensor.shape} " f"and requested shape {shape}" ) op = TensorBroadcastTo(shape, dtype=tensor.dtype, sparse=tensor.issparse()) return op(tensor, shape)