maxframe 1.3.0__cp311-cp311-macosx_10_9_universal2.whl → 2.0.0b1__cp311-cp311-macosx_10_9_universal2.whl

maxframe/tensor/indexing/choose.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,9 +12,12 @@
 # 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
@@ -40,10 +41,11 @@ class TensorChoose(TensorOperator, TensorOperatorMixin):
 
         super().__setattr__(key, value)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self.a = self._inputs[0]
-        self.choices = self._inputs[1:]
+    @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):

maxframe/tensor/indexing/compress.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

maxframe/tensor/indexing/extract.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

maxframe/tensor/indexing/fill_diagonal.py

@@ -12,10 +12,12 @@
 # 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
+from ...core import ENTITY_TYPE, EntityData
 from ...serialization.serializables import AnyField, BoolField, Int32Field, KeyField
 from ..core import Tensor
 from ..datasource import tensor as astensor
@@ -35,11 +37,12 @@ class TensorFillDiagonal(TensorOperator, TensorOperatorMixin):
     def input(self):
         return self._input
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self._input = self._inputs[0]
-        if len(self._inputs) == 2:
-            self.val = self._inputs[1]
+    @classmethod
+    def _set_inputs(cls, op: "TensorFillDiagonal", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op._input = op._inputs[0]
+        if len(op._inputs) == 2:
+            op.val = op._inputs[1]
 
     def __call__(self, a, val=None):
         inputs = [a]

maxframe/tensor/indexing/flatnonzero.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -52,7 +50,7 @@ def flatnonzero(a):
     Use the indices of the non-zero elements as an index array to extract
     these elements:
 
-    >>> x.ravel()[mt.flatnonzero(x)].execute()  # TODO(jisheng): accomplish this after fancy indexing is supported
+    >>> x.ravel()[mt.flatnonzero(x)].execute()
 
     """
     from ..misc import ravel

maxframe/tensor/indexing/getitem.py

@@ -13,20 +13,15 @@
 # limitations under the License.
 
 from numbers import Integral
+from typing import List
 
 import numpy as np
 
 from ... import opcodes
-from ...core import ENTITY_TYPE
-from ...serialization.serializables import (
-    FieldTypes,
-    Int32Field,
-    KeyField,
-    ListField,
-    TupleField,
-)
+from ...core import ENTITY_TYPE, EntityData
+from ...serialization.serializables import KeyField, ListField
 from ..core import TENSOR_TYPE, TensorOrder
-from ..operators import TensorHasInput, TensorMapReduceOperator, TensorOperatorMixin
+from ..operators import TensorHasInput, TensorOperatorMixin
 from ..utils import filter_inputs
 from .core import calc_shape, process_index
 
@@ -39,14 +34,15 @@ class TensorIndex(TensorHasInput, TensorOperatorMixin):
     _input = KeyField("input")
     indexes = ListField("indexes", default=None)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        inputs_iter = iter(self._inputs[1:])
+    @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 self.indexes
+            for index in op.indexes
         ]
-        self.indexes = new_indexes
+        op.indexes = new_indexes
 
     def on_output_modify(self, new_output):
         from .setitem import TensorIndexSetValue
@@ -72,35 +68,6 @@ class TensorIndex(TensorHasInput, TensorOperatorMixin):
         return self.new_tensor(filter_inputs([a] + list(index)), shape, order=order)
 
 
-class FancyIndexingDistribute(TensorMapReduceOperator, TensorOperatorMixin):
-    _op_type_ = opcodes.FANCY_INDEX_DISTRIBUTE
-
-    _input = KeyField("input")
-    dest_nsplits = TupleField(
-        "dest_nsplits", FieldTypes.tuple(FieldTypes.uint64), default=None
-    )
-    axes = TupleField("axes", FieldTypes.int32, default=None)
-
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self._input = self._inputs[0]
-
-    @property
-    def output_limit(self):
-        return len(self.axes) + 1
-
-
-class FancyIndexingConcat(TensorMapReduceOperator, TensorOperatorMixin):
-    _op_type_ = opcodes.FANCY_INDEX_CONCAT
-
-    fancy_index_axis = Int32Field("fancy_index_axis", default=None)
-    fancy_index_shape = TupleField("fancy_index_shape", FieldTypes.int64, default=None)
-
-    @property
-    def input(self):
-        return self._input
-
-
 def _is_bool_index(index_obj):
     return isinstance(index_obj, TENSOR_TYPE) and index_obj.dtype == np.bool_
 

maxframe/tensor/indexing/nonzero.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -92,9 +90,9 @@ def nonzero(a):
     >>> mt.nonzero(x).execute()
     (array([0, 1, 2, 2]), array([0, 1, 0, 1]))
 
-    >>> x[mt.nonzero(x)].execute()  # TODO(jisheng): accomplish this after fancy indexing is supported
+    >>> x[mt.nonzero(x)].execute()
 
-    >>> mt.transpose(mt.nonzero(x)).execute() # TODO(jisheng): accomplish this later
+    >>> 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

maxframe/tensor/indexing/setitem.py

@@ -13,11 +13,13 @@
 # limitations under the License.
 
 from numbers import Integral
+from typing import List
 
 import numpy as np
 
 from ... import opcodes
-from ...core import ENTITY_TYPE
+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
@@ -57,17 +59,25 @@ class TensorIndexSetValue(TensorMapReduceOperator, TensorOperatorMixin):
             **kw,
         )
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self.input = self._inputs[0]
-        inputs_iter = iter(self._inputs[1:])
+    @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 self.indexes
+            for index in op.indexes
         ]
-        self.indexes = tuple(new_indexes)
-        if isinstance(self.value, ENTITY_TYPE):
-            self.value = next(inputs_iter)
+        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])

maxframe/tensor/indexing/slice.py

@@ -12,8 +12,10 @@
 # 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
 
@@ -24,6 +26,7 @@ class TensorSlice(TensorHasInput, TensorOperatorMixin):
     _input = KeyField("input")
     slices = ListField("slices", default=None)
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self._input = self._inputs[0]
+    @classmethod
+    def _set_inputs(cls, op: "TensorSlice", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op._input = op._inputs[0]

maxframe/tensor/indexing/take.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

maxframe/tensor/indexing/tests/__init__.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

maxframe/tensor/indexing/tests/test_indexing.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

maxframe/tensor/indexing/unravel_index.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright 1999-2025 Alibaba Group Holding Ltd.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,11 +13,12 @@
 # limitations under the License.
 
 from collections.abc import Iterable
+from typing import List
 
 import numpy as np
 
 from ... import opcodes
-from ...core import ExecutableTuple
+from ...core import EntityData, ExecutableTuple
 from ...serialization.serializables import FieldTypes, KeyField, StringField, TupleField
 from ..core import TensorOrder
 from ..datasource import tensor as astensor
@@ -42,9 +41,10 @@ class TensorUnravelIndex(TensorHasInput, TensorOperatorMixin):
     def output_limit(self):
         return float("inf")
 
-    def _set_inputs(self, inputs):
-        super()._set_inputs(inputs)
-        self._input = self._inputs[0]
+    @classmethod
+    def _set_inputs(cls, op: "TensorUnravelIndex", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op._input = op._inputs[0]
 
     def __call__(self, indices):
         order = TensorOrder.C_ORDER if self.order == "C" else TensorOrder.F_ORDER

maxframe/tensor/lib/__init__.py

@@ -0,0 +1,16 @@
+# 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 . import index_tricks
+from .index_tricks import nd_grid