maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/tensor/indexing/__init__.py

@@ -0,0 +1,47 @@
+# 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 .choose import TensorChoose, choose
+from .compress import compress
+from .extract import extract
+from .fill_diagonal import TensorFillDiagonal, fill_diagonal
+from .flatnonzero import flatnonzero
+from .getitem import TensorIndex
+from .nonzero import TensorNonzero, nonzero
+from .setitem import TensorIndexSetValue
+from .slice import TensorSlice
+from .take import take
+from .unravel_index import TensorUnravelIndex, unravel_index
+
+
+def _install():
+    from ..core import Tensor, TensorData
+    from .getitem import _getitem
+    from .setitem import _setitem
+
+    setattr(Tensor, "__getitem__", _getitem)
+    setattr(TensorData, "__getitem__", _getitem)
+    setattr(Tensor, "__setitem__", _setitem)
+    setattr(Tensor, "take", take)
+    setattr(
+        Tensor,
+        "compress",
+        lambda a, condition, axis=None: compress(condition, a, axis=axis),
+    )
+    setattr(Tensor, "choose", choose)
+    setattr(Tensor, "nonzero", nonzero)
+
+
+_install()
+del _install

maxframe/tensor/indexing/choose.py

@@ -0,0 +1,198 @@
+# 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 EntityData
+from ...serialization.serializables import FieldTypes, KeyField, ListField, StringField
+from ..core import Tensor, TensorOrder
+from ..datasource import tensor as astensor
+from ..operators import TensorOperator, TensorOperatorMixin
+from ..utils import broadcast_shape, check_out_param
+
+
+class TensorChoose(TensorOperator, TensorOperatorMixin):
+    _op_type_ = opcodes.CHOOSE
+
+    a = KeyField("a", default=None)
+    choices = ListField("choices", FieldTypes.key, default=None)
+    mode = StringField("mode", default=None)
+
+    def __init__(self, mode=None, **kw):
+        super().__init__(mode=mode, **kw)
+
+    def __setattr__(self, key, value):
+        if key == "mode" and value not in ("raise", "wrap", "clip"):
+            raise ValueError(f"mode should be raise, wrap or clip, not {value}")
+
+        super().__setattr__(key, value)
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorChoose", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.a = op._inputs[0]
+        op.choices = op._inputs[1:]
+
+    def __call__(self, a, choices, out=None):
+        if out is not None and not isinstance(out, Tensor):
+            raise TypeError(f"out should be Tensor object, got {type(out)} instead")
+
+        inputs = [a] + choices
+        shape = broadcast_shape(a.shape, *[c.shape for c in choices])
+        order = TensorOrder.C_ORDER if out is None else out.order
+        t = self.new_tensor(inputs, shape, order=order)
+
+        if out is None:
+            return t
+
+        check_out_param(out, t, "unsafe")
+        out_shape, out_dtype = out.shape, out.dtype
+        # if `out` is specified, use out's dtype and shape
+        if out_shape != t.shape:
+            raise ValueError(f"output shape should be {t.shape}, got {out_shape}")
+        setattr(self, "dtype", out_dtype)
+        out.data = t.data
+        return out
+
+
+def choose(a, choices, out=None, mode="raise"):
+    """
+    Construct a tensor from an index tensor and a set of tensors to choose from.
+
+    First of all, if confused or uncertain, definitely look at the Examples -
+    in its full generality, this function is less simple than it might
+    seem from the following code description (below ndi =
+    `mt.lib.index_tricks`):
+
+    ``mt.choose(a,c) == mt.array([c[a[I]][I] for I in ndi.ndindex(a.shape)])``.
+
+    But this omits some subtleties.  Here is a fully general summary:
+
+    Given an "index" tensor (`a`) of integers and a sequence of `n` tensors
+    (`choices`), `a` and each choice tensor are first broadcast, as necessary,
+    to tensors of a common shape; calling these *Ba* and *Bchoices[i], i =
+    0,...,n-1* we have that, necessarily, ``Ba.shape == Bchoices[i].shape``
+    for each `i`.  Then, a new array with shape ``Ba.shape`` is created as
+    follows:
+
+    * if ``mode=raise`` (the default), then, first of all, each element of
+      `a` (and thus `Ba`) must be in the range `[0, n-1]`; now, suppose that
+      `i` (in that range) is the value at the `(j0, j1, ..., jm)` position
+      in `Ba` - then the value at the same position in the new array is the
+      value in `Bchoices[i]` at that same position;
+
+    * if ``mode=wrap``, values in `a` (and thus `Ba`) may be any (signed)
+      integer; modular arithmetic is used to map integers outside the range
+      `[0, n-1]` back into that range; and then the new array is constructed
+      as above;
+
+    * if ``mode=clip``, values in `a` (and thus `Ba`) may be any (signed)
+      integer; negative integers are mapped to 0; values greater than `n-1`
+      are mapped to `n-1`; and then the new tensor is constructed as above.
+
+    Parameters
+    ----------
+    a : int tensor
+        This tensor must contain integers in `[0, n-1]`, where `n` is the number
+        of choices, unless ``mode=wrap`` or ``mode=clip``, in which cases any
+        integers are permissible.
+    choices : sequence of tensors
+        Choice tensors. `a` and all of the choices must be broadcastable to the
+        same shape.  If `choices` is itself a tensor (not recommended), then
+        its outermost dimension (i.e., the one corresponding to
+        ``choices.shape[0]``) is taken as defining the "sequence".
+    out : tensor, optional
+        If provided, the result will be inserted into this tensor. It should
+        be of the appropriate shape and dtype.
+    mode : {'raise' (default), 'wrap', 'clip'}, optional
+        Specifies how indices outside `[0, n-1]` will be treated:
+
+          * 'raise' : an exception is raised
+          * 'wrap' : value becomes value mod `n`
+          * 'clip' : values < 0 are mapped to 0, values > n-1 are mapped to n-1
+
+    Returns
+    -------
+    merged_array : Tensor
+        The merged result.
+
+    Raises
+    ------
+    ValueError: shape mismatch
+        If `a` and each choice tensor are not all broadcastable to the same
+        shape.
+
+    See Also
+    --------
+    Tensor.choose : equivalent method
+
+    Notes
+    -----
+    To reduce the chance of misinterpretation, even though the following
+    "abuse" is nominally supported, `choices` should neither be, nor be
+    thought of as, a single tensor, i.e., the outermost sequence-like container
+    should be either a list or a tuple.
+
+    Examples
+    --------
+
+    >>> import maxframe.tensor as mt
+
+    >>> choices = [[0, 1, 2, 3], [10, 11, 12, 13],
+    ...   [20, 21, 22, 23], [30, 31, 32, 33]]
+    >>> mt.choose([2, 3, 1, 0], choices
+    ... # the first element of the result will be the first element of the
+    ... # third (2+1) "array" in choices, namely, 20; the second element
+    ... # will be the second element of the fourth (3+1) choice array, i.e.,
+    ... # 31, etc.
+    ... ).execute()
+    array([20, 31, 12,  3])
+    >>> mt.choose([2, 4, 1, 0], choices, mode='clip').execute() # 4 goes to 3 (4-1)
+    array([20, 31, 12,  3])
+    >>> # because there are 4 choice arrays
+    >>> mt.choose([2, 4, 1, 0], choices, mode='wrap').execute() # 4 goes to (4 mod 4)
+    array([20,  1, 12,  3])
+    >>> # i.e., 0
+
+    A couple examples illustrating how choose broadcasts:
+
+    >>> a = [[1, 0, 1], [0, 1, 0], [1, 0, 1]]
+    >>> choices = [-10, 10]
+    >>> mt.choose(a, choices).execute()
+    array([[ 10, -10,  10],
+           [-10,  10, -10],
+           [ 10, -10,  10]])
+
+    >>> # With thanks to Anne Archibald
+    >>> a = mt.array([0, 1]).reshape((2,1,1))
+    >>> c1 = mt.array([1, 2, 3]).reshape((1,3,1))
+    >>> c2 = mt.array([-1, -2, -3, -4, -5]).reshape((1,1,5))
+    >>> mt.choose(a, (c1, c2)).execute() # result is 2x3x5, res[0,:,:]=c1, res[1,:,:]=c2
+    array([[[ 1,  1,  1,  1,  1],
+            [ 2,  2,  2,  2,  2],
+            [ 3,  3,  3,  3,  3]],
+           [[-1, -2, -3, -4, -5],
+            [-1, -2, -3, -4, -5],
+            [-1, -2, -3, -4, -5]]])
+
+    """
+    a = astensor(a, dtype="i8")
+    choices = [astensor(c) for c in choices]
+
+    dtype = np.result_type(*[c.dtype for c in choices])
+    op = TensorChoose(mode=mode, dtype=dtype)
+    return op(a, choices, out=out)

maxframe/tensor/indexing/compress.py

@@ -0,0 +1,122 @@
+# 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 ..core import Tensor
+from ..datasource import tensor as astensor
+from ..utils import validate_axis
+
+
+def compress(condition, a, axis=None, out=None):
+    """
+    Return selected slices of a tensor along given axis.
+
+    When working along a given axis, a slice along that axis is returned in
+    `output` for each index where `condition` evaluates to True. When
+    working on a 1-D array, `compress` is equivalent to `extract`.
+
+    Parameters
+    ----------
+    condition : 1-D tensor of bools
+        Tensor that selects which entries to return. If len(condition)
+        is less than the size of `a` along the given axis, then output is
+        truncated to the length of the condition tensor.
+    a : array_like
+        Tensor from which to extract a part.
+    axis : int, optional
+        Axis along which to take slices. If None (default), work on the
+        flattened tensor.
+    out : Tensor, optional
+        Output tensor.  Its type is preserved and it must be of the right
+        shape to hold the output.
+
+    Returns
+    -------
+    compressed_array : Tensor
+        A copy of `a` without the slices along axis for which `condition`
+        is false.
+
+    See Also
+    --------
+    take, choose, diag, diagonal, select
+    Tensor.compress : Equivalent method in ndarray
+    mt.extract: Equivalent method when working on 1-D arrays
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> a = mt.array([[1, 2], [3, 4], [5, 6]])
+    >>> a.execute()
+    array([[1, 2],
+           [3, 4],
+           [5, 6]])
+    >>> mt.compress([0, 1], a, axis=0).execute()
+    array([[3, 4]])
+    >>> mt.compress([False, True, True], a, axis=0).execute()
+    array([[3, 4],
+           [5, 6]])
+    >>> mt.compress([False, True], a, axis=1).execute()
+    array([[2],
+           [4],
+           [6]])
+
+    Working on the flattened tensor does not return slices along an axis but
+    selects elements.
+
+    >>> mt.compress([False, True], a).execute()
+    array([2])
+
+    """
+    a = astensor(a)
+    condition = astensor(condition, dtype=bool)
+
+    if condition.ndim != 1:
+        raise ValueError("condition must be an 1-d tensor")
+
+    if axis is None:
+        a = a.ravel()
+        if len(condition) < a.size:
+            a = a[: len(condition)]
+        return a[condition]
+
+    try:
+        axis = validate_axis(a.ndim, axis)
+    except ValueError:
+        raise np.AxisError(
+            f"axis {axis} is out of bounds for tensor of dimension {a.ndim}"
+        )
+
+    try:
+        if len(condition) < a.shape[axis]:
+            a = a[(slice(None),) * axis + (slice(len(condition)),)]
+        t = a[(slice(None),) * axis + (condition,)]
+        if out is None:
+            return t
+
+        if out is not None and not isinstance(out, Tensor):
+            raise TypeError(f"out should be Tensor object, got {type(out)} instead")
+        if not np.can_cast(out.dtype, t.dtype, "safe"):
+            raise TypeError(
+                f"Cannot cast array data from dtype('{out.dtype}') to dtype('{t.dtype}') "
+                "according to the rule 'safe'"
+            )
+        # skip shape check because out shape is unknown
+        out.data = t.astype(out.dtype, order=out.order.value).data
+        return out
+    except IndexError:
+        raise np.AxisError(
+            f"axis {len(condition)} is out of bounds for tensor of dimension 1"
+        )

maxframe/tensor/indexing/core.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.
+
+
+import itertools
+from numbers import Integral
+
+import numpy as np
+
+from ..core import TENSOR_TYPE
+from ..datasource import tensor as astensor
+from ..utils import broadcast_shape, calc_sliced_size, index_ndim, replace_ellipsis
+
+_INDEX_ERROR_MSG = (
+    "only integers, slices (`:`), ellipsis (`...`), "
+    "numpy.newaxis (`None`) and integer or boolean arrays are valid indices"
+)
+
+
+def calc_shape(tensor_shape, index):
+    shape = []
+    in_axis = 0
+    out_axis = 0
+    fancy_index = None
+    fancy_index_shapes = []
+    for ind in index:
+        if isinstance(ind, TENSOR_TYPE + (np.ndarray,)) and ind.dtype == np.bool_:
+            # bool
+            shape.append(np.nan if not isinstance(ind, np.ndarray) else int(ind.sum()))
+            for i, t_size, size in zip(
+                itertools.count(0),
+                ind.shape,
+                tensor_shape[in_axis : ind.ndim + in_axis],
+            ):
+                if not np.isnan(t_size) and not np.isnan(size) and t_size != size:
+                    raise IndexError(
+                        f"boolean index did not match indexed array along dimension {in_axis + i}; "
+                        f"dimension is {size} but corresponding boolean dimension is {t_size}"
+                    )
+            in_axis += ind.ndim
+            out_axis += 1
+        elif isinstance(ind, TENSOR_TYPE + (np.ndarray,)):
+            first_fancy_index = False
+            if fancy_index is None:
+                first_fancy_index = True
+                fancy_index = out_axis
+            if isinstance(ind, np.ndarray) and np.any(ind >= tensor_shape[in_axis]):
+                out_of_range_index = next(
+                    i for i in ind.flat if i >= tensor_shape[in_axis]
+                )
+                raise IndexError(
+                    f"IndexError: index {out_of_range_index} is out of "
+                    f"bounds with size {tensor_shape[in_axis]}"
+                )
+            fancy_index_shapes.append(ind.shape)
+            in_axis += 1
+            if first_fancy_index:
+                out_axis += ind.ndim
+        elif isinstance(ind, slice):
+            if np.isnan(tensor_shape[in_axis]):
+                shape.append(np.nan)
+            else:
+                shape.append(calc_sliced_size(tensor_shape[in_axis], ind))
+            in_axis += 1
+            out_axis += 1
+        elif isinstance(ind, Integral):
+            size = tensor_shape[in_axis]
+            if not np.isnan(size) and ind >= size:
+                raise IndexError(
+                    f"index {ind} is out of bounds for axis {in_axis} with size {size}"
+                )
+            in_axis += 1
+        else:
+            assert ind is None
+            shape.append(1)
+
+    if fancy_index is not None:
+        try:
+            if any(np.isnan(np.prod(s)) for s in fancy_index_shapes):
+                fancy_index_shape = (np.nan,) * len(fancy_index_shapes[0])
+            else:
+                fancy_index_shape = broadcast_shape(*fancy_index_shapes)
+            shape = shape[:fancy_index] + list(fancy_index_shape) + shape[fancy_index:]
+        except ValueError:
+            raise IndexError(
+                "shape mismatch: indexing arrays could not be broadcast together "
+                "with shapes {0}".format(" ".join(str(s) for s in fancy_index_shapes))
+            )
+
+    return shape
+
+
+def preprocess_index(index, convert_bool_to_fancy=None):
+    from .nonzero import nonzero
+
+    inds = []
+    fancy_indexes = []
+    bool_indexes = []
+    all_fancy_index_ndarray = True
+    all_bool_index_ndarray = True
+    for j, ind in enumerate(index):
+        if isinstance(ind, (list, np.ndarray) + TENSOR_TYPE):
+            if not isinstance(ind, TENSOR_TYPE):
+                ind = np.array(ind)
+            if ind.dtype.kind not in "biu":
+                raise IndexError(_INDEX_ERROR_MSG)
+            if ind.dtype.kind == "b":
+                # bool indexing
+                bool_indexes.append(j)
+                if not isinstance(ind, np.ndarray):
+                    all_bool_index_ndarray = False
+            else:
+                # fancy indexing
+                fancy_indexes.append(j)
+                if not isinstance(ind, np.ndarray):
+                    all_fancy_index_ndarray = False
+        elif (
+            not isinstance(ind, (slice, Integral))
+            and ind is not None
+            and ind is not Ellipsis
+        ):
+            raise IndexError(_INDEX_ERROR_MSG)
+        inds.append(ind)
+
+    if convert_bool_to_fancy is None:
+        convert_bool_to_fancy = (fancy_indexes and len(bool_indexes) > 0) or len(
+            bool_indexes
+        ) > 1
+
+    if not all_fancy_index_ndarray or (
+        convert_bool_to_fancy and not all_bool_index_ndarray
+    ):
+        # if not all fancy indexes are ndarray,
+        # or bool indexes need to be converted to fancy indexes,
+        # and not all bool indexes are ndarray,
+        # we will convert all of them to Tensor
+        for fancy_index in fancy_indexes:
+            inds[fancy_index] = astensor(inds[fancy_index])
+
+    # convert bool index to fancy index when any situation below meets:
+    # 1. fancy indexes and bool indexes both exists
+    # 2. bool indexes more than 2
+    if convert_bool_to_fancy:
+        default_m = None
+        if len(fancy_indexes) > 0:
+            default_m = (
+                np.nonzero
+                if isinstance(inds[fancy_indexes[0]], np.ndarray)
+                else nonzero
+            )
+        for bool_index in bool_indexes:
+            ind = inds[bool_index]
+            m = default_m
+            if m is None:
+                m = np.nonzero if isinstance(ind, np.ndarray) else nonzero
+            ind = m(ind)[0]
+            inds[bool_index] = ind
+
+    return tuple(inds)
+
+
+def process_index(tensor_ndim, item, convert_bool_to_fancy=None):
+    if isinstance(item, list):
+        arr = np.array(item)
+        if arr.dtype == object:
+            item = tuple(item)
+        elif arr.dtype.kind == "f":
+            raise IndexError(_INDEX_ERROR_MSG)
+        else:
+            item = (arr,)
+    elif not isinstance(item, tuple):
+        item = (item,)
+
+    index = preprocess_index(item, convert_bool_to_fancy=convert_bool_to_fancy)
+    index = replace_ellipsis(index, tensor_ndim)
+    missing = tensor_ndim - sum(index_ndim(i) for i in index)
+    if missing < 0:
+        raise IndexError("too many indices for tensor")
+    return index + (slice(None),) * missing

maxframe/tensor/indexing/extract.py

@@ -0,0 +1,69 @@
+# 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 ..datasource import tensor as astensor
+
+
+def extract(condition, a):
+    """
+    Return the elements of a tensor that satisfy some condition.
+
+    This is equivalent to ``mt.compress(ravel(condition), ravel(arr))``.  If
+    `condition` is boolean ``mt.extract`` is equivalent to ``arr[condition]``.
+
+    Note that `place` does the exact opposite of `extract`.
+
+    Parameters
+    ----------
+    condition : array_like
+        An array whose nonzero or True entries indicate the elements of `arr`
+        to extract.
+    a : array_like
+        Input tensor of the same size as `condition`.
+
+    Returns
+    -------
+    extract : Tensor
+        Rank 1 tensor of values from `arr` where `condition` is True.
+
+    See Also
+    --------
+    take, put, copyto, compress, place
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> arr = mt.arange(12).reshape((3, 4))
+    >>> arr.execute()
+    array([[ 0,  1,  2,  3],
+           [ 4,  5,  6,  7],
+           [ 8,  9, 10, 11]])
+    >>> condition = mt.mod(arr, 3)==0
+    >>> condition.execute()
+    array([[ True, False, False,  True],
+           [False, False,  True, False],
+           [False,  True, False, False]])
+    >>> mt.extract(condition, arr).execute()
+    array([0, 3, 6, 9])
+
+
+    If `condition` is boolean:
+
+    >>> arr[condition].execute()
+    array([0, 3, 6, 9])
+
+    """
+    condition = astensor(condition, dtype=bool)
+    return a[condition]