maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/tensor/indexing/fill_diagonal.py

@@ -0,0 +1,180 @@
+# 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 ENTITY_TYPE, EntityData
+from ...serialization.serializables import AnyField, BoolField, Int32Field
+from ..core import Tensor
+from ..datasource import tensor as astensor
+from ..operators import TensorOperator, TensorOperatorMixin
+
+
+class TensorFillDiagonal(TensorOperator, TensorOperatorMixin):
+    _op_type_ = opcodes.FILL_DIAGONAL
+
+    val = AnyField("val", default=None)
+    wrap = BoolField("wrap", default=None)
+    # used for chunk
+    k = Int32Field("k", default=None)
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorFillDiagonal", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        if len(op._inputs) == 2:
+            op.val = op._inputs[1]
+
+    def __call__(self, a, val=None):
+        inputs = [a]
+        if val is not None:
+            inputs.append(val)
+        return self.new_tensor(inputs, shape=a.shape, order=a.order)
+
+
+def fill_diagonal(a, val, wrap=False):
+    """Fill the main diagonal of the given tensor of any dimensionality.
+
+    For a tensor `a` with ``a.ndim >= 2``, the diagonal is the list of
+    locations with indices ``a[i, ..., i]`` all identical. This function
+    modifies the input tensor in-place, it does not return a value.
+
+    Parameters
+    ----------
+    a : Tensor, at least 2-D.
+      Tensor whose diagonal is to be filled, it gets modified in-place.
+
+    val : scalar
+      Value to be written on the diagonal, its type must be compatible with
+      that of the tensor a.
+
+    wrap : bool
+      For tall matrices in NumPy version up to 1.6.2, the
+      diagonal "wrapped" after N columns. You can have this behavior
+      with this option. This affects only tall matrices.
+
+    See also
+    --------
+    diag_indices, diag_indices_from
+
+    Notes
+    -----
+
+    This functionality can be obtained via `diag_indices`, but internally
+    this version uses a much faster implementation that never constructs the
+    indices and uses simple slicing.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> a = mt.zeros((3, 3), int)
+    >>> mt.fill_diagonal(a, 5)
+    >>> a.execute()
+    array([[5, 0, 0],
+           [0, 5, 0],
+           [0, 0, 5]])
+
+    The same function can operate on a 4-D tensor:
+
+    >>> a = mt.zeros((3, 3, 3, 3), int)
+    >>> mt.fill_diagonal(a, 4)
+
+    We only show a few blocks for clarity:
+
+    >>> a[0, 0].execute()
+    array([[4, 0, 0],
+           [0, 0, 0],
+           [0, 0, 0]])
+    >>> a[1, 1].execute()
+    array([[0, 0, 0],
+           [0, 4, 0],
+           [0, 0, 0]])
+    >>> a[2, 2].execute()
+    array([[0, 0, 0],
+           [0, 0, 0],
+           [0, 0, 4]])
+
+    The wrap option affects only tall matrices:
+
+    >>> # tall matrices no wrap
+    >>> a = mt.zeros((5, 3), int)
+    >>> mt.fill_diagonal(a, 4)
+    >>> a.execute()
+    array([[4, 0, 0],
+           [0, 4, 0],
+           [0, 0, 4],
+           [0, 0, 0],
+           [0, 0, 0]])
+
+    >>> # tall matrices wrap
+    >>> a = mt.zeros((5, 3), int)
+    >>> mt.fill_diagonal(a, 4, wrap=True)
+    >>> a.execute()
+    array([[4, 0, 0],
+           [0, 4, 0],
+           [0, 0, 4],
+           [0, 0, 0],
+           [4, 0, 0]])
+
+    >>> # wide matrices
+    >>> a = mt.zeros((3, 5), int)
+    >>> mt.fill_diagonal(a, 4, wrap=True)
+    >>> a.execute()
+    array([[4, 0, 0, 0, 0],
+           [0, 4, 0, 0, 0],
+           [0, 0, 4, 0, 0]])
+
+    The anti-diagonal can be filled by reversing the order of elements
+    using either `numpy.flipud` or `numpy.fliplr`.
+
+    >>> a = mt.zeros((3, 3), int)
+    >>> mt.fill_diagonal(mt.fliplr(a), [1,2,3])  # Horizontal flip
+    >>> a.execute()
+    array([[0, 0, 1],
+           [0, 2, 0],
+           [3, 0, 0]])
+    >>> mt.fill_diagonal(mt.flipud(a), [1,2,3])  # Vertical flip
+    >>> a.execute()
+    array([[0, 0, 3],
+           [0, 2, 0],
+           [1, 0, 0]])
+
+    Note that the order in which the diagonal is filled varies depending
+    on the flip function.
+    """
+
+    if not isinstance(a, Tensor):
+        raise TypeError(f"`a` should be a tensor, got {type(a)}")
+    if a.ndim < 2:
+        raise ValueError("array must be at least 2-d")
+    if a.ndim > 2 and len(set(a.shape)) != 1:
+        raise ValueError("All dimensions of input must be of equal length")
+
+    # process val
+    if isinstance(val, ENTITY_TYPE):
+        val = astensor(val)
+        if val.ndim > 1:
+            val = val.ravel()
+        val_input = val
+    else:
+        val = np.asarray(val)
+        if val.ndim > 1:
+            val = val.ravel()
+        val_input = None
+
+    op = TensorFillDiagonal(val=val, wrap=wrap, dtype=a.dtype)
+    t = op(a, val=val_input)
+    a.data = t.data

maxframe/tensor/indexing/flatnonzero.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.
+
+from .nonzero import nonzero
+
+
+def flatnonzero(a):
+    """
+    Return indices that are non-zero in the flattened version of a.
+
+    This is equivalent to a.ravel().nonzero()[0].
+
+    Parameters
+    ----------
+    a : Tensor
+        Input tensor.
+
+    Returns
+    -------
+    res : Tensor
+        Output tensor, containing the indices of the elements of `a.ravel()`
+        that are non-zero.
+
+    See Also
+    --------
+    nonzero : Return the indices of the non-zero elements of the input tensor.
+    ravel : Return a 1-D tensor containing the elements of the input tensor.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> x = mt.arange(-2, 3)
+    >>> x.execute()
+    array([-2, -1,  0,  1,  2])
+    >>> mt.flatnonzero(x).execute()
+    array([0, 1, 3, 4])
+
+    Use the indices of the non-zero elements as an index array to extract
+    these elements:
+
+    >>> x.ravel()[mt.flatnonzero(x)].execute()
+
+    """
+    from ..misc import ravel
+
+    return nonzero(ravel(a))[0]

maxframe/tensor/indexing/getitem.py

@@ -0,0 +1,144 @@
+# 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 numbers import Integral
+from typing import List
+
+import numpy as np
+
+from ... import opcodes
+from ...core import ENTITY_TYPE, EntityData
+from ...serialization.serializables import KeyField, ListField
+from ..core import TENSOR_TYPE, TensorOrder
+from ..operators import TensorHasInput, TensorOperatorMixin
+from ..utils import filter_inputs
+from .core import calc_shape, process_index
+
+FANCY_INDEX_TYPES = TENSOR_TYPE + (np.ndarray,)
+
+
+class TensorIndex(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.INDEX
+
+    _input = KeyField("input")
+    indexes = ListField("indexes", default=None)
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorIndex", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        inputs_iter = iter(op._inputs[1:])
+        new_indexes = [
+            next(inputs_iter) if isinstance(index, ENTITY_TYPE) else index
+            for index in op.indexes
+        ]
+        op.indexes = new_indexes
+
+    def on_output_modify(self, new_output):
+        from .setitem import TensorIndexSetValue
+
+        if self.create_view:
+            a = self.input
+            op = TensorIndexSetValue(
+                dtype=a.dtype,
+                sparse=a.issparse(),
+                indexes=tuple(self.indexes),
+                value=new_output,
+            )
+            return op(a, self.indexes, new_output)
+
+    def on_input_modify(self, new_input):
+        if self.create_view:
+            new_op = self.copy().reset_key()
+            new_inputs = [new_input] + self.inputs[1:]
+            return new_op.new_tensor(new_inputs, shape=self.outputs[0].shape)
+
+    def __call__(self, a, index, shape, order):
+        self.indexes = list(index)
+        return self.new_tensor(filter_inputs([a] + list(index)), shape, order=order)
+
+
+def _is_bool_index(index_obj):
+    return isinstance(index_obj, TENSOR_TYPE) and index_obj.dtype == np.bool_
+
+
+def _is_fancy_index(index_obj):
+    return isinstance(index_obj, FANCY_INDEX_TYPES) and index_obj.dtype != np.bool_
+
+
+def _is_create_view(index):
+    # is view if all of index is slice, int or newaxis
+    return all(isinstance(ind, (slice, Integral)) or ind is None for ind in index)
+
+
+def _calc_order(a, index):
+    if a.order == TensorOrder.C_ORDER:
+        return TensorOrder.C_ORDER
+
+    in_axis = 0
+    for ind in index:
+        if _is_bool_index(ind):
+            in_axis += ind.ndim
+            return TensorOrder.C_ORDER
+        elif _is_fancy_index(ind):
+            in_axis += 1
+            return TensorOrder.C_ORDER
+        elif ind is None:
+            continue
+        elif isinstance(ind, slice):
+            shape = a.shape[in_axis]
+            if shape is np.nan:
+                return TensorOrder.C_ORDER
+            slc = ind.indices(shape)
+            if slc[0] == 0 and slc[1] == shape and slc[2] == 1:
+                continue
+            else:
+                return TensorOrder.C_ORDER
+        else:
+            assert isinstance(ind, Integral)
+            in_axis += 1
+            return TensorOrder.C_ORDER
+
+    return TensorOrder.F_ORDER
+
+
+def _getitem_nocheck(a, item, convert_bool_to_fancy=None):
+    index = process_index(a.ndim, item, convert_bool_to_fancy=convert_bool_to_fancy)
+    if convert_bool_to_fancy is False:
+        # come from __setitem__, the bool index is not converted to fancy index
+        # if multiple bool indexes or bool + fancy indexes exist,
+        # thus the shape will be wrong,
+        # here we just convert when calculating shape,
+        # refer to issue #1282.
+        shape = calc_shape(a.shape, process_index(a.ndim, index))
+    else:
+        shape = calc_shape(a.shape, index)
+    tensor_order = _calc_order(a, index)
+    op = TensorIndex(
+        dtype=a.dtype,
+        sparse=a.issparse(),
+        indexes=list(index),
+        create_view=_is_create_view(index),
+    )
+    return op(a, index, tuple(shape), order=tensor_order)
+
+
+def _getitem(a, item):
+    if isinstance(item, (list, tuple)) and all(
+        isinstance(it, slice) and it == slice(None) for it in item
+    ):
+        # nothing to do
+        return a
+
+    # TODO(jisheng): field access, e.g. t['a'], t[['a', 'b']]
+    return _getitem_nocheck(a, item)

maxframe/tensor/indexing/nonzero.py

@@ -0,0 +1,118 @@
+# 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 KeyField
+from ..core import TensorOrder
+from ..datasource import tensor as astensor
+from ..operators import TensorHasInput, TensorOperatorMixin
+
+
+class TensorNonzero(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.NONZERO
+
+    _input = KeyField("input")
+
+    @property
+    def output_limit(self):
+        return float("inf")
+
+    def __call__(self, a):
+        kws = [
+            {"shape": (np.nan,), "order": TensorOrder.C_ORDER, "_idx_": i}
+            for i in range(a.ndim)
+        ]
+        return ExecutableTuple(self.new_tensors([a], kws=kws, output_limit=len(kws)))
+
+
+def nonzero(a):
+    """
+    Return the indices of the elements that are non-zero.
+
+    Returns a tuple of tensors, one for each dimension of `a`,
+    containing the indices of the non-zero elements in that
+    dimension. The values in `a` are always tested and returned.
+    The corresponding non-zero
+    values can be obtained with::
+
+        a[nonzero(a)]
+
+    To group the indices by element, rather than dimension, use::
+
+        transpose(nonzero(a))
+
+    The result of this is always a 2-D array, with a row for
+    each non-zero element.
+
+    Parameters
+    ----------
+    a : array_like
+        Input tensor.
+
+    Returns
+    -------
+    tuple_of_arrays : tuple
+        Indices of elements that are non-zero.
+
+    See Also
+    --------
+    flatnonzero :
+        Return indices that are non-zero in the flattened version of the input
+        tensor.
+    Tensor.nonzero :
+        Equivalent tensor method.
+    count_nonzero :
+        Counts the number of non-zero elements in the input tensor.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> x = mt.array([[1,0,0], [0,2,0], [1,1,0]])
+    >>> x.execute()
+    array([[1, 0, 0],
+           [0, 2, 0],
+           [1, 1, 0]])
+    >>> mt.nonzero(x).execute()
+    (array([0, 1, 2, 2]), array([0, 1, 0, 1]))
+
+    >>> x[mt.nonzero(x)].execute()
+
+    >>> mt.transpose(mt.nonzero(x)).execute()
+
+    A common use for ``nonzero`` is to find the indices of an array, where
+    a condition is True.  Given an array `a`, the condition `a` > 3 is a
+    boolean array and since False is interpreted as 0, np.nonzero(a > 3)
+    yields the indices of the `a` where the condition is true.
+
+    >>> a = mt.array([[1,2,3],[4,5,6],[7,8,9]])
+    >>> (a > 3).execute()
+    array([[False, False, False],
+           [ True,  True,  True],
+           [ True,  True,  True]])
+    >>> mt.nonzero(a > 3).execute()
+    (array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))
+
+    The ``nonzero`` method of the boolean array can also be called.
+
+    >>> (a > 3).nonzero().execute()
+    (array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))
+
+    """
+    a = astensor(a)
+    op = TensorNonzero(dtype=np.dtype(np.intp))
+    return op(a)

maxframe/tensor/indexing/setitem.py

@@ -0,0 +1,142 @@
+# 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 numbers import Integral
+from typing import List
+
+import numpy as np
+
+from ... import opcodes
+from ...core import ENTITY_TYPE, EntityData
+from ...core.operator import OperatorStage
+from ...serialization.serializables import AnyField, BoolField, KeyField, TupleField
+from ...tensor import tensor as astensor
+from ..core import TENSOR_TYPE
+from ..operators import TensorMapReduceOperator, TensorOperatorMixin
+from ..utils import filter_inputs
+from .core import process_index
+
+
+class TensorIndexSetValue(TensorMapReduceOperator, TensorOperatorMixin):
+    _op_type_ = opcodes.INDEXSETVALUE
+
+    input = KeyField("input")
+    indexes = TupleField("indexes")
+    value = AnyField("value")
+    is_fancy_index = BoolField("is_fancy_index")
+    input_nsplits = TupleField("input_nsplits")
+    chunk_offsets = TupleField("chunk_offsets")
+    shuffle_axes = TupleField("shuffle_axes")
+
+    def __init__(
+        self,
+        indexes=None,
+        value=None,
+        is_fancy_index=None,
+        input_nsplits=None,
+        chunk_offsets=None,
+        shuffle_axes=None,
+        **kw,
+    ):
+        super().__init__(
+            indexes=indexes,
+            value=value,
+            is_fancy_index=is_fancy_index,
+            input_nsplits=input_nsplits,
+            chunk_offsets=chunk_offsets,
+            shuffle_axes=shuffle_axes,
+            **kw,
+        )
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorIndexSetValue", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.input = op._inputs[0]
+        if op.stage == OperatorStage.reduce:
+            op.input = op._inputs[0]
+            return
+        elif op.stage == OperatorStage.map:
+            inputs_iter = iter(op._inputs)
+        else:
+            op.input = op._inputs[0]
+            inputs_iter = iter(op._inputs[1:])
+        new_indexes = [
+            next(inputs_iter) if isinstance(index, ENTITY_TYPE) else index
+            for index in op.indexes
+        ]
+        op.indexes = tuple(new_indexes)
+        if isinstance(op.value, ENTITY_TYPE):
+            op.value = next(inputs_iter)
+
+    def __call__(self, a, index, value):
+        inputs = filter_inputs([a] + list(index) + [value])
+        self.indexes = tuple(index)
+        self.value = value
+        return self.new_tensor(inputs, a.shape, order=a.order)
+
+    def on_output_modify(self, new_output):
+        return new_output
+
+    def on_input_modify(self, new_input):
+        new_op = self.copy().reset_key()
+        new_inputs = [new_input] + self.inputs[1:]
+        return new_op.new_tensor(new_inputs, shape=self.outputs[0].shape)
+
+
+def _check_support(indexes):
+    if all(
+        (
+            isinstance(ix, (TENSOR_TYPE, np.ndarray))
+            and ix.dtype != np.bool_
+            or isinstance(ix, slice)
+            and ix == slice(None)
+        )
+        for ix in indexes
+    ):
+        if any(isinstance(ix, (TENSOR_TYPE, np.ndarray)) for ix in indexes):
+            return True
+    for index in indexes:
+        if isinstance(index, (slice, Integral)):
+            pass
+        elif isinstance(index, (np.ndarray, TENSOR_TYPE)) and index.dtype == np.bool_:
+            pass
+        else:  # pragma: no cover
+            raise NotImplementedError(
+                "Only slice, int, or bool indexing "
+                f"supported by now, got {type(index)}"
+            )
+    return False
+
+
+def _setitem(a, item, value):
+    index = process_index(a.ndim, item, convert_bool_to_fancy=False)
+    if not (np.isscalar(value) or (isinstance(value, tuple) and a.dtype.fields)):
+        # do not convert for tuple when dtype is record type.
+        value = astensor(value)
+
+    is_fancy_index = _check_support(index)
+    if is_fancy_index:
+        index = [astensor(ind) if isinstance(ind, np.ndarray) else ind for ind in index]
+
+    # __setitem__ on a view should be still a view, see GH #732.
+    op = TensorIndexSetValue(
+        dtype=a.dtype,
+        sparse=a.issparse(),
+        is_fancy_index=is_fancy_index,
+        indexes=tuple(index),
+        value=value,
+        create_view=a.op.create_view,
+    )
+    ret = op(a, index, value)
+    a.data = ret.data

maxframe/tensor/indexing/slice.py

@@ -0,0 +1,32 @@
+# 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
+
+from ... import opcodes
+from ...core import EntityData
+from ...serialization.serializables import KeyField, ListField
+from ..operators import TensorHasInput, TensorOperatorMixin
+
+
+class TensorSlice(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.SLICE
+
+    _input = KeyField("input")
+    slices = ListField("slices", default=None)
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorSlice", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op._input = op._inputs[0]