maxframe 2.2.0__cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

maxframe/tensor/misc/searchsorted.py

@@ -0,0 +1,147 @@
+# Copyright 1999-2025 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.
+
+from typing import List
+
+import numpy as np
+
+from ... import opcodes
+from ...core import TILEABLE_TYPE, EntityData
+from ...serialization.serializables import AnyField, Int32Field, Int64Field, StringField
+from ..core import TENSOR_TYPE, TensorOrder
+from ..datasource.array import tensor as astensor
+from ..operators import TensorOperator, TensorOperatorMixin
+
+
+class TensorSearchsorted(TensorOperator, TensorOperatorMixin):
+    _op_type_ = opcodes.SEARCHSORTED
+
+    a = AnyField("a")
+    v = AnyField("v")
+    side = StringField("side")
+    combine_size = Int32Field("combine_size")
+    # for chunk
+    offset = Int64Field("offset")
+    size = Int64Field("size")
+    n_chunk = Int64Field("n_chunk")
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorSearchsorted", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.a = inputs[0]
+        if isinstance(op.v, TILEABLE_TYPE):
+            op.v = op._inputs[1]
+
+    def __call__(self, a, v):
+        inputs = [a]
+        if isinstance(v, TILEABLE_TYPE):
+            inputs.append(v)
+            shape = v.shape
+        else:
+            shape = ()
+        return self.new_tensor(inputs, shape=shape, order=TensorOrder.C_ORDER)
+
+
+def searchsorted(a, v, side="left", sorter=None, combine_size=None):
+    """
+    Find indices where elements should be inserted to maintain order.
+
+    Find the indices into a sorted tensor `a` such that, if the
+    corresponding elements in `v` were inserted before the indices, the
+    order of `a` would be preserved.
+
+    Assuming that `a` is sorted:
+
+    ======  ============================
+    `side`  returned index `i` satisfies
+    ======  ============================
+    left    ``a[i-1] < v <= a[i]``
+    right   ``a[i-1] <= v < a[i]``
+    ======  ============================
+
+    Parameters
+    ----------
+    a : 1-D array_like
+        Input tensor. If `sorter` is None, then it must be sorted in
+        ascending order, otherwise `sorter` must be an array of indices
+        that sort it.
+    v : array_like
+        Values to insert into `a`.
+    side : {'left', 'right'}, optional
+        If 'left', the index of the first suitable location found is given.
+        If 'right', return the last such index.  If there is no suitable
+        index, return either 0 or N (where N is the length of `a`).
+    sorter : 1-D array_like, optional
+        Optional tensor of integer indices that sort array a into ascending
+        order. They are typically the result of argsort.
+    combine_size: int, optional
+        The number of chunks to combine.
+
+    Returns
+    -------
+    indices : tensor of ints
+        Array of insertion points with the same shape as `v`.
+
+    See Also
+    --------
+    sort : Return a sorted copy of a tensor.
+    histogram : Produce histogram from 1-D data.
+
+    Notes
+    -----
+    Binary search is used to find the required insertion points.
+
+    This function is a faster version of the builtin python `bisect.bisect_left`
+    (``side='left'``) and `bisect.bisect_right` (``side='right'``) functions,
+    which is also vectorized in the `v` argument.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> mt.searchsorted([1,2,3,4,5], 3).execute()
+    2
+    >>> mt.searchsorted([1,2,3,4,5], 3, side='right').execute()
+    3
+    >>> mt.searchsorted([1,2,3,4,5], [-10, 10, 2, 3]).execute()
+    array([0, 5, 1, 2])
+
+    """
+
+    if (
+        not isinstance(a, TENSOR_TYPE)
+        and sorter is not None
+        and not isinstance(sorter, TENSOR_TYPE)
+    ):
+        a = astensor(np.asarray(a)[sorter])
+    else:
+        a = astensor(a)
+        if sorter is not None:
+            a = a[sorter]
+
+    if a.ndim != 1:
+        raise ValueError("`a` should be 1-d tensor")
+    if a.issparse():
+        # does not support sparse tensor
+        raise ValueError("`a` should be a dense tensor")
+    if side not in {"left", "right"}:
+        raise ValueError(f"'{side}' is an invalid value for keyword 'side'")
+
+    if not np.isscalar(v):
+        v = astensor(v)
+
+    op = TensorSearchsorted(
+        v=v, side=side, dtype=np.dtype(np.intp), combine_size=combine_size
+    )
+    return op(a, v)

maxframe/tensor/misc/setdiff1d.py

@@ -0,0 +1,58 @@
+# Copyright 1999-2025 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.
+
+
+def setdiff1d(ar1, ar2, assume_unique=False):
+    """
+    Find the set difference of two tensors.
+
+    Return the unique values in `ar1` that are not in `ar2`.
+
+    Parameters
+    ----------
+    ar1 : array_like
+        Input tensor.
+    ar2 : array_like
+        Input comparison tensor.
+    assume_unique : bool
+        If True, the input tensors are both assumed to be unique, which
+        can speed up the calculation.  Default is False.
+
+    Returns
+    -------
+    setdiff1d : Tensor
+        1D tensor of values in `ar1` that are not in `ar2`. The result
+        is sorted when `assume_unique=False`, but otherwise only sorted
+        if the input is sorted.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> a = mt.array([1, 2, 3, 2, 4, 1])
+    >>> b = mt.array([3, 4, 5, 6])
+    >>> mt.setdiff1d(a, b).execute()
+    array([1, 2])
+
+    """
+
+    from ..datasource.array import asarray
+    from .in1d import in1d
+    from .unique import unique
+
+    if assume_unique:
+        ar1 = asarray(ar1).ravel()
+    else:
+        ar1 = unique(ar1)
+        ar2 = unique(ar2)
+    return ar1[in1d(ar1, ar2, assume_unique=True, invert=True)]

maxframe/tensor/misc/shape.py

@@ -0,0 +1,89 @@
+# Copyright 1999-2025 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
+from ..core import TensorOrder
+from ..datasource import tensor as astensor
+from ..operators import TensorOperator, TensorOperatorMixin
+
+
+class TensorGetShape(TensorOperator, TensorOperatorMixin):
+    _op_type_ = opcodes.GET_SHAPE
+
+    ndim = Int32Field("ndim")
+
+    def __init__(self, pure_depends=None, **kw):
+        super().__init__(_pure_depends=pure_depends, **kw)
+
+    @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])
+
+        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))
+
+
+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 maxframe.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)

maxframe/tensor/misc/split.py

@@ -0,0 +1,190 @@
+# Copyright 1999-2025 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.
+
+from typing import List
+
+import numpy as np
+
+from ... import opcodes
+from ...core import ExecutableTuple
+from ...lib.sparse.core import get_array_module
+from ...serialization.serializables import AnyField, Int32Field
+from ...typing_ import EntityType
+from ..core import Tensor
+from ..datasource import tensor as astensor
+from ..operators import TensorHasInput, TensorOperatorMixin
+from ..utils import calc_sliced_size
+
+
+class TensorSplit(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.ARRAY_SPLIT
+
+    indices_or_sections = AnyField("indices_or_sections", default=None)
+    axis = Int32Field("axis", default=None)
+
+    @property
+    def output_limit(self):
+        return float("inf")
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorSplit", inputs: List[EntityType]):
+        super()._set_inputs(op, inputs)
+        if len(inputs) > 1:
+            op.indices_or_sections = inputs[1]
+
+    def __call__(self, a, indices_or_sections, is_split=False):
+        axis = self.axis
+        size = a.shape[axis]
+        if np.isnan(size):
+            raise ValueError(
+                "cannot split array with unknown shape, "
+                "call `.execute()` on input tensor first"
+            )
+
+        if (
+            isinstance(indices_or_sections, Tensor)
+            and hasattr(indices_or_sections.op, "data")
+            and indices_or_sections.op.data is not None
+        ):
+            indices_or_sections = indices_or_sections.op.data
+
+        try:
+            indices_or_sections = int(indices_or_sections)
+            if is_split:
+                if size % indices_or_sections:
+                    raise ValueError(
+                        "tensor split does not result in an equal division"
+                    )
+                nparts = indices_or_sections
+                nsplit = (size // indices_or_sections,) * nparts
+            else:
+                nparts = indices_or_sections
+                if size % indices_or_sections == 0:
+                    nsplit = (size // indices_or_sections,) * nparts
+                else:
+                    nsplit = (size // indices_or_sections + 1,) * (
+                        size % indices_or_sections
+                    ) + (size // indices_or_sections,) * (
+                        size - size % indices_or_sections
+                    )
+        except TypeError:
+            if isinstance(indices_or_sections, Tensor):
+                nparts = indices_or_sections.shape[0] + 1
+                nsplit = (np.nan,) * nparts
+            else:
+                ind = indices_or_sections = get_array_module(
+                    indices_or_sections
+                ).asarray(indices_or_sections)
+                if indices_or_sections.ndim != 1 or not np.issubdtype(
+                    indices_or_sections.dtype, np.integer
+                ):
+                    raise TypeError("slice indices must be integers or None")
+                nparts = indices_or_sections.shape[0] + 1
+                get = lambda i: None if i < 0 or i >= len(ind) else ind[i]
+                nsplit = [
+                    calc_sliced_size(size, slice(get(j - 1), get(j)))
+                    for j in range(nparts)
+                ]
+
+        inputs = [a]
+        if isinstance(indices_or_sections, Tensor):
+            inputs.append(indices_or_sections)
+        else:
+            self.indices_or_sections = indices_or_sections
+
+        kws = [
+            {
+                "i": i,
+                "shape": a.shape[:axis] + (nsplit[i],) + a.shape[axis + 1 :],
+                "order": a.order,
+            }
+            for i in range(nparts)
+        ]
+        return ExecutableTuple(self.new_tensors(inputs, kws=kws, output_limit=nparts))
+
+
+def _split(a, indices_or_sections, axis=0, is_split=False):
+    op = TensorSplit(axis=axis, dtype=a.dtype)
+    return op(a, indices_or_sections, is_split=is_split)
+
+
+def split(ary, indices_or_sections, axis=0):
+    """
+    Split a tensor into multiple sub-tensors.
+
+    Parameters
+    ----------
+    ary : Tensor
+        Tensor to be divided into sub-tensors.
+    indices_or_sections : int or 1-D tensor
+        If `indices_or_sections` is an integer, N, the array will be divided
+        into N equal tensors along `axis`.  If such a split is not possible,
+        an error is raised.
+
+        If `indices_or_sections` is a 1-D tensor of sorted integers, the entries
+        indicate where along `axis` the array is split.  For example,
+        ``[2, 3]`` would, for ``axis=0``, result in
+
+          - ary[:2]
+          - ary[2:3]
+          - ary[3:]
+
+        If an index exceeds the dimension of the tensor along `axis`,
+        an empty sub-tensor is returned correspondingly.
+    axis : int, optional
+        The axis along which to split, default is 0.
+
+    Returns
+    -------
+    sub-tensors : list of Tensors
+        A list of sub-tensors.
+
+    Raises
+    ------
+    ValueError
+        If `indices_or_sections` is given as an integer, but
+        a split does not result in equal division.
+
+    See Also
+    --------
+    array_split : Split a tensor into multiple sub-tensors of equal or
+                  near-equal size.  Does not raise an exception if
+                  an equal division cannot be made.
+    hsplit : Split  into multiple sub-arrays horizontally (column-wise).
+    vsplit : Split tensor into multiple sub-tensors vertically (row wise).
+    dsplit : Split tensor into multiple sub-tensors along the 3rd axis (depth).
+    concatenate : Join a sequence of tensors along an existing axis.
+    stack : Join a sequence of tensors along a new axis.
+    hstack : Stack tensors in sequence horizontally (column wise).
+    vstack : Stack tensors in sequence vertically (row wise).
+    dstack : Stack tensors in sequence depth wise (along third dimension).
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> x = mt.arange(9.0)
+    >>> mt.split(x, 3).execute()
+    [array([ 0.,  1.,  2.]), array([ 3.,  4.,  5.]), array([ 6.,  7.,  8.])]
+
+    >>> x = mt.arange(8.0)
+    >>> mt.split(x, [3, 5, 6, 10]).execute()
+    [array([ 0.,  1.,  2.]),
+     array([ 3.,  4.]),
+     array([ 5.]),
+     array([ 6.,  7.]),
+     array([], dtype=float64)]
+
+    """
+    return _split(astensor(ary), indices_or_sections, axis=axis, is_split=True)

maxframe/tensor/misc/squeeze.py

@@ -0,0 +1,117 @@
+# Copyright 1999-2025 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.
+
+from collections.abc import Iterable
+
+from ... import opcodes
+from ...serialization.serializables import FieldTypes, TupleField
+from ..operators import TensorHasInput, TensorOperatorMixin
+
+
+class TensorSqueeze(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.SQUEEZE
+
+    axis = TupleField("axis", FieldTypes.int32)
+
+    def on_output_modify(self, new_output):
+        slcs = [slice(None)] * new_output.ndim
+        for axis in self.axis:
+            slcs.insert(axis, None)
+        return new_output[slcs]
+
+    def on_input_modify(self, new_input):
+        op = self.copy().reset_key()
+        return op(new_input, self.outputs[0].shape)
+
+    @staticmethod
+    def _get_squeeze_shape(shape, axis):
+        if axis is not None:
+            if isinstance(axis, Iterable):
+                axis = tuple(axis)
+            else:
+                axis = (axis,)
+
+            for ax in axis:
+                if shape[ax] != 1:
+                    raise ValueError(
+                        "cannot select an axis to squeeze out "
+                        "which has size not equal to one"
+                    )
+            shape = tuple(s for i, s in enumerate(shape) if i not in axis)
+        else:
+            axis = tuple(i for i, s in enumerate(shape) if s == 1)
+            shape = tuple(s for s in shape if s != 1)
+
+        return shape, axis
+
+    def __call__(self, a, shape):
+        return self.new_tensor([a], shape, order=a.order)
+
+
+def squeeze(a, axis=None):
+    """
+    Remove single-dimensional entries from the shape of a tensor.
+
+    Parameters
+    ----------
+    a : array_like
+        Input data.
+    axis : None or int or tuple of ints, optional
+        Selects a subset of the single-dimensional entries in the
+        shape. If an axis is selected with shape entry greater than
+        one, an error is raised.
+
+    Returns
+    -------
+    squeezed : Tensor
+        The input tensor, but with all or a subset of the
+        dimensions of length 1 removed. This is always `a` itself
+        or a view into `a`.
+
+    Raises
+    ------
+    ValueError
+        If `axis` is not `None`, and an axis being squeezed is not of length 1
+
+    See Also
+    --------
+    expand_dims : The inverse operation, adding singleton dimensions
+    reshape : Insert, remove, and combine dimensions, and resize existing ones
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> x = mt.array([[[0], [1], [2]]])
+    >>> x.shape
+    (1, 3, 1)
+    >>> mt.squeeze(x).shape
+    (3,)
+    >>> mt.squeeze(x, axis=0).shape
+    (3, 1)
+    >>> mt.squeeze(x, axis=1).shape
+    Traceback (most recent call last):
+    ...
+    ValueError: cannot select an axis to squeeze out which has size not equal to one
+    >>> mt.squeeze(x, axis=2).shape
+    (1, 3)
+
+    """
+    shape, axis = TensorSqueeze._get_squeeze_shape(a.shape, axis)
+
+    if 1 not in a.shape:
+        return a
+
+    op = TensorSqueeze(axis=axis, dtype=a.dtype, sparse=a.issparse())
+    return op(a, shape)