# 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
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 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)
@property
def shape(self):
return self._shape
def __call__(self, tensor, shape):
return self.new_tensor([tensor], shape)
@classmethod
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
)
out_chunks.append(out_chunk)
nsplits = [
tuple(
c.shape[i]
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(
[in_tensor],
tensor.shape,
order=tensor.order,
chunks=out_chunks,
nsplits=nsplits,
)
@classmethod
def execute(cls, ctx, op):
xp = get_array_module(ctx[op.input.key])
input_data = ctx[op.input.key]
device_id = input_data.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)