maxframe 2.0.0b2__cp37-cp37m-win32.whl → 2.3.0rc1__cp37-cp37m-win32.whl

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/tile.py

@@ -0,0 +1,109 @@
+# 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
+
+
+def tile(A, reps):
+    """
+    Construct a tensor by repeating A the number of times given by reps.
+
+    If `reps` has length ``d``, the result will have dimension of
+    ``max(d, A.ndim)``.
+
+    If ``A.ndim < d``, `A` is promoted to be d-dimensional by prepending new
+    axes. So a shape (3,) array is promoted to (1, 3) for 2-D replication,
+    or shape (1, 1, 3) for 3-D replication. If this is not the desired
+    behavior, promote `A` to d-dimensions manually before calling this
+    function.
+
+    If ``A.ndim > d``, `reps` is promoted to `A`.ndim by prepending 1's to it.
+    Thus for an `A` of shape (2, 3, 4, 5), a `reps` of (2, 2) is treated as
+    (1, 1, 2, 2).
+
+    Note : Although tile may be used for broadcasting, it is strongly
+    recommended to use MaxFrame's broadcasting operations and functions.
+
+    Parameters
+    ----------
+    A : array_like
+        The input tensor.
+    reps : array_like
+        The number of repetitions of `A` along each axis.
+
+    Returns
+    -------
+    c : Tensor
+        The tiled output tensor.
+
+    See Also
+    --------
+    repeat : Repeat elements of a tensor.
+    broadcast_to : Broadcast a tensor to a new shape
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> a = mt.array([0, 1, 2])
+    >>> mt.tile(a, 2).execute()
+    array([0, 1, 2, 0, 1, 2])
+    >>> mt.tile(a, (2, 2)).execute()
+    array([[0, 1, 2, 0, 1, 2],
+           [0, 1, 2, 0, 1, 2]])
+    >>> mt.tile(a, (2, 1, 2)).execute()
+    array([[[0, 1, 2, 0, 1, 2]],
+           [[0, 1, 2, 0, 1, 2]]])
+
+    >>> b = mt.array([[1, 2], [3, 4]])
+    >>> mt.tile(b, 2).execute()
+    array([[1, 2, 1, 2],
+           [3, 4, 3, 4]])
+    >>> mt.tile(b, (2, 1)).execute()
+    array([[1, 2],
+           [3, 4],
+           [1, 2],
+           [3, 4]])
+
+    >>> c = mt.array([1,2,3,4])
+    >>> mt.tile(c,(4,1)).execute()
+    array([[1, 2, 3, 4],
+           [1, 2, 3, 4],
+           [1, 2, 3, 4],
+           [1, 2, 3, 4]])
+    """
+    from ..merge import concatenate
+
+    try:
+        tup = tuple(reps)
+    except TypeError:
+        tup = (reps,)
+
+    d = len(tup)
+    if A.ndim < d:
+        A = A[tuple(np.newaxis for _ in range(d - A.ndim))]
+    elif A.ndim > d:
+        tup = (1,) * (A.ndim - d) + tup
+
+    a = A
+    for axis, rep in enumerate(tup):
+        if rep == 0:
+            slc = (slice(None),) * axis + (slice(0),)
+            a = a[slc]
+        elif rep < 0:
+            raise ValueError("negative dimensions are not allowed")
+        elif rep > 1:
+            a = concatenate([a] * rep, axis=axis)
+
+    return a

maxframe/tensor/misc/vsplit.py

@@ -0,0 +1,74 @@
+# 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 .split import split
+
+
+def vsplit(a, indices_or_sections):
+    """
+    Split a tensor into multiple sub-tensors vertically (row-wise).
+
+    Please refer to the ``split`` documentation.  ``vsplit`` is equivalent
+    to ``split`` with `axis=0` (default), the tensor is always split along the
+    first axis regardless of the tensor dimension.
+
+    See Also
+    --------
+    split : Split a tensor into multiple sub-tensors of equal size.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> x = mt.arange(16.0).reshape(4, 4)
+    >>> x.execute()
+    array([[  0.,   1.,   2.,   3.],
+           [  4.,   5.,   6.,   7.],
+           [  8.,   9.,  10.,  11.],
+           [ 12.,  13.,  14.,  15.]])
+    >>> mt.vsplit(x, 2).execute()
+    [array([[ 0.,  1.,  2.,  3.],
+           [ 4.,  5.,  6.,  7.]]),
+     array([[  8.,   9.,  10.,  11.],
+           [ 12.,  13.,  14.,  15.]])]
+    >>> mt.vsplit(x, mt.array([3, 6])).execute()
+    [array([[  0.,   1.,   2.,   3.],
+           [  4.,   5.,   6.,   7.],
+           [  8.,   9.,  10.,  11.]]),
+     array([[ 12.,  13.,  14.,  15.]]),
+     array([], dtype=float64)]
+
+    With a higher dimensional tensor the split is still along the first axis.
+
+    >>> x = mt.arange(8.0).reshape(2, 2, 2)
+    >>> x.execute()
+    array([[[ 0.,  1.],
+            [ 2.,  3.]],
+           [[ 4.,  5.],
+            [ 6.,  7.]]])
+    >>> mt.vsplit(x, 2).execute()
+    [array([[[ 0.,  1.],
+            [ 2.,  3.]]]),
+     array([[[ 4.,  5.],
+            [ 6.,  7.]]])]
+
+    """
+    from ..datasource import tensor as astensor
+
+    ary = a
+    a = astensor(a)
+
+    if a.ndim < 2:
+        raise ValueError("vsplit only works on tensors of 2 or more dimensions")
+    return split(ary, indices_or_sections, 0)

maxframe/tensor/reduction/array_equal.py

@@ -48,6 +48,7 @@ def array_equal(a1, a2):
     False
 
     """
+    from ..arithmetic import equal
     from ..datasource import tensor as astensor
     from ..datasource.scalar import scalar
     from .all import all
@@ -59,4 +60,4 @@ def array_equal(a1, a2):
 
     if a1.shape != a2.shape:
         return scalar(False)
-    return all(astensor(a1 == a2))
+    return all(astensor(equal(a1, a2, magic=True)))

maxframe/tensor/sort/__init__.py

@@ -12,5 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .argpartition import argpartition
 from .argsort import argsort
+from .partition import partition
 from .sort import sort

maxframe/tensor/sort/argpartition.py

@@ -0,0 +1,98 @@
+# 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 .partition import TensorPartition, _validate_partition_arguments
+
+
+def argpartition(a, kth, axis=-1, kind="introselect", order=None, **kw):
+    """
+    Perform an indirect partition along the given axis using the
+    algorithm specified by the `kind` keyword. It returns an array of
+    indices of the same shape as `a` that index data along the given
+    axis in partitioned order.
+
+    .. versionadded:: 1.8.0
+
+    Parameters
+    ----------
+    a : array_like
+        Tensor to sort.
+    kth : int or sequence of ints
+        Element index to partition by. The k-th element will be in its
+        final sorted position and all smaller elements will be moved
+        before it and all larger elements behind it. The order all
+        elements in the partitions is undefined. If provided with a
+        sequence of k-th it will partition all of them into their sorted
+        position at once.
+    axis : int or None, optional
+        Axis along which to sort. The default is -1 (the last axis). If
+        None, the flattened tensor is used.
+    kind : {'introselect'}, optional
+        Selection algorithm. Default is 'introselect'
+    order : str or list of str, optional
+        When `a` is a tensor with fields defined, this argument
+        specifies which fields to compare first, second, etc. A single
+        field can be specified as a string, and not all fields need be
+        specified, but unspecified fields will still be used, in the
+        order in which they come up in the dtype, to break ties.
+
+    Returns
+    -------
+    index_tensor : Tensor, int
+        Tensor of indices that partition `a` along the specified axis.
+        If `a` is one-dimensional, ``a[index_tensor]`` yields a partitioned `a`.
+        More generally, ``np.take_along_axis(a, index_tensor, axis=a)`` always
+        yields the partitioned `a`, irrespective of dimensionality.
+
+    See Also
+    --------
+    partition : Describes partition algorithms used.
+    Tensor.partition : Inplace partition.
+    argsort : Full indirect sort
+
+    Notes
+    -----
+    See `partition` for notes on the different selection algorithms.
+
+    Examples
+    --------
+    One dimensional tensor:
+
+    >>> import maxframe.tensor as mt
+    >>> x = mt.array([3, 4, 2, 1])
+    >>> x[mt.argpartition(x, 3)].execute()
+    array([2, 1, 3, 4])
+    >>> x[mt.argpartition(x, (1, 3))].execute()
+    array([1, 2, 3, 4])
+
+    >>> x = [3, 4, 2, 1]
+    >>> mt.array(x)[mt.argpartition(x, 3)].execute()
+    array([2, 1, 3, 4])
+
+    """
+    a, kth, axis, kind, order, need_align = _validate_partition_arguments(
+        a, kth, axis, kind, order, kw
+    )
+    op = TensorPartition(
+        kth=kth,
+        axis=axis,
+        kind=kind,
+        order=order,
+        need_align=need_align,
+        return_value=False,
+        return_indices=True,
+        dtype=a.dtype,
+        gpu=a.op.gpu,
+    )
+    return op(a, kth)