# Copyright 2022-2023 XProbe Inc.
# derived from 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
from ...core import ExecutableTuple
from ...serialization.serializables import Int32Field, KeyField
from ...utils import calc_nsplits
from ..core import TensorOrder
from ..datasource import tensor as astensor
from ..operands import TensorOperand, TensorOperandMixin
class TensorGetShape(TensorOperand, TensorOperandMixin):
_op_type_ = opcodes.GET_SHAPE
_a = KeyField("a")
_ndim = Int32Field("ndim")
def __init__(self, pure_depends=None, a=None, ndim=None, dtype=None, **kw):
super().__init__(
dtype=dtype, _a=a, _ndim=ndim, _pure_depends=pure_depends, **kw
)
@property
def a(self):
return self._a
@property
def ndim(self):
return self._ndim
def _set_inputs(self, inputs):
super()._set_inputs(inputs)
if self._a is not None:
self._a = self._inputs[0]
@property
def output_limit(self):
return self._ndim
def __call__(self, a):
if not np.isnan(a.size):
return ExecutableTuple([astensor(s) for s in a.shape])
self._a = a
kws = []
for i in range(self.output_limit):
kws.append(
{
"shape": (),
"dtype": np.dtype(np.intc),
"order": TensorOrder.C_ORDER,
"i": i,
}
)
return ExecutableTuple(self.new_tensors([a], kws=kws))
@classmethod
def tile(cls, op):
a = op.a
outs = op.outputs
yield a.chunks
chunk_op = TensorGetShape(pure_depends=[True] * len(a.chunks), ndim=op.ndim)
chunk_kws = []
for out in outs:
params = out.params
params["index"] = ()
chunk_kws.append(params)
chunks = chunk_op.new_chunks(a.chunks, kws=chunk_kws)
kws = []
for c, out in zip(chunks, outs):
params = out.params
params["chunks"] = [c]
params["nsplits"] = ()
kws.append(params)
new_op = op.copy()
return new_op.new_tensors(op.inputs, kws=kws, output_limit=op.output_limit)
@classmethod
def execute(cls, ctx, op):
chunk_idx_to_chunk_shapes = dict((c.index, c.shape) for c in op.inputs)
nsplits = calc_nsplits(chunk_idx_to_chunk_shapes)
shape = tuple(sum(ns) for ns in nsplits)
for o, s in zip(op.outputs, shape):
ctx[o.key] = s
[docs]def shape(a):
"""
Return the shape of a tensor.
Parameters
----------
a : array_like
Input tensor.
Returns
-------
shape : ExecutableTuple of tensors
The elements of the shape tuple give the lengths of the
corresponding array dimensions.
Examples
--------
>>> import mars.tensor as mt
>>> mt.shape(mt.eye(3)).execute()
(3, 3)
>>> mt.shape([[1, 2]]).execute()
(1, 2)
>>> mt.shape([0]).execute()
(1,)
>>> mt.shape(0).execute()
()
>>> a = mt.array([(1, 2), (3, 4)], dtype=[('x', 'i4'), ('y', 'i4')])
>>> mt.shape(a).execute()
(2,)
"""
a = astensor(a)
op = TensorGetShape(ndim=a.ndim)
return op(a)