maxframe 2.4.0rc1__cp312-cp312-win32.whl

maxframe/tensor/datasource/zeros.py

@@ -0,0 +1,192 @@
+# 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 ...serialization.serializables import (
+    AnyField,
+    FieldTypes,
+    KeyField,
+    StringField,
+    TupleField,
+)
+from ...utils import on_deserialize_shape, on_serialize_shape
+from ..utils import get_order
+from .array import tensor
+from .core import TensorLike, TensorNoInput
+
+
+class TensorZeros(TensorNoInput):
+    _op_type_ = opcodes.TENSOR_ZEROS
+
+    order = StringField("order")
+    shape = TupleField(
+        "shape",
+        FieldTypes.int64,
+        on_serialize=on_serialize_shape,
+        on_deserialize=on_deserialize_shape,
+    )
+    chunk_size = AnyField("chunk_size")
+
+    def __init__(self, shape=None, **kwargs):
+        if type(shape) is int:
+            shape = (shape,)
+        super().__init__(shape=shape, **kwargs)
+
+
+def zeros(shape, dtype=None, chunk_size=None, gpu=None, sparse=False, order="C"):
+    """
+    Return a new tensor of given shape and type, filled with zeros.
+
+    Parameters
+    ----------
+    shape : int or sequence of ints
+        Shape of the new tensor, e.g., ``(2, 3)`` or ``2``.
+    dtype : data-type, optional
+        The desired data-type for the array, e.g., `mt.int8`.  Default is
+        `mt.float64`.
+    chunk_size : int or tuple of int or tuple of ints, optional
+        Desired chunk size on each dimension
+    gpu : bool, optional
+        Allocate the tensor on GPU if True, False as default
+    sparse: bool, optional
+        Create sparse tensor if True, False as default
+    order : {'C', 'F'}, optional, default: 'C'
+        Whether to store multi-dimensional data in row-major
+        (C-style) or column-major (Fortran-style) order in
+        memory.
+
+    Returns
+    -------
+    out : Tensor
+        Tensor of zeros with the given shape, dtype, and order.
+
+    See Also
+    --------
+    zeros_like : Return a tensor of zeros with shape and type of input.
+    ones_like : Return a tensor of ones with shape and type of input.
+    empty_like : Return a empty tensor with shape and type of input.
+    ones : Return a new tensor setting values to one.
+    empty : Return a new uninitialized tensor.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> mt.zeros(5).execute()
+    array([ 0.,  0.,  0.,  0.,  0.])
+
+    >>> mt.zeros((5,), dtype=int).execute()
+    array([0, 0, 0, 0, 0])
+
+    >>> mt.zeros((2, 1)).execute()
+    array([[ 0.],
+           [ 0.]])
+
+    >>> s = (2,2)
+    >>> mt.zeros(s).execute()
+    array([[ 0.,  0.],
+           [ 0.,  0.]])
+
+    >>> mt.zeros((2,), dtype=[('x', 'i4'), ('y', 'i4')]).execute() # custom dtype
+    array([(0, 0), (0, 0)],
+          dtype=[('x', '<i4'), ('y', '<i4')])
+    """
+    tensor_order = get_order(
+        order,
+        None,
+        available_options="CF",
+        err_msg="only 'C' or 'F' order is permitted",
+    )
+    dtype = np.dtype(dtype or "f8")
+    op = TensorZeros(
+        dtype=dtype,
+        shape=shape,
+        chunk_size=chunk_size,
+        gpu=gpu,
+        sparse=sparse,
+        order=order,
+    )
+    return op(shape, chunk_size=chunk_size, order=tensor_order)
+
+
+class TensorZerosLike(TensorLike):
+    _op_type_ = opcodes.TENSOR_ZEROS_LIKE
+
+    _input = KeyField("input")
+    order = StringField("order")
+
+    def __init__(self, dtype=None, gpu=None, sparse=False, **kw):
+        dtype = np.dtype(dtype) if dtype is not None else None
+        super().__init__(dtype=dtype, gpu=gpu, sparse=sparse, **kw)
+
+
+def zeros_like(a, dtype=None, gpu=None, order="K"):
+    """
+    Return a tensor of zeros with the same shape and type as a given tensor.
+
+    Parameters
+    ----------
+    a : array_like
+        The shape and data-type of `a` define these same attributes of
+        the returned array.
+    dtype : data-type, optional
+        Overrides the data type of the result.
+    gpu : bool, optional
+        Allocate the tensor on GPU if True, None as default
+    order : {'C', 'F', 'A', or 'K'}, optional
+        Overrides the memory layout of the result. 'C' means C-order,
+        'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous,
+        'C' otherwise. 'K' means match the layout of `a` as closely
+        as possible.
+
+    Returns
+    -------
+    out : Tensor
+        tensor of zeros with the same shape and type as `a`.
+
+    See Also
+    --------
+    ones_like : Return an array of ones with shape and type of input.
+    empty_like : Return an empty array with shape and type of input.
+    zeros : Return a new array setting values to zero.
+    ones : Return a new array setting values to one.
+    empty : Return a new uninitialized array.
+
+    Examples
+    --------
+    >>> import maxframe.tensr as mt
+    >>> x = mt.arange(6)
+    >>> x = x.reshape((2, 3))
+    >>> x.execute()
+    array([[0, 1, 2],
+           [3, 4, 5]])
+
+    >>> mt.zeros_like(x).execute()
+    array([[0, 0, 0],
+           [0, 0, 0]])
+
+    >>> y = mt.arange(3, dtype=float)
+    >>> y.execute()
+    array([ 0.,  1.,  2.])
+
+    >>> mt.zeros_like(y).execute()
+    array([ 0.,  0.,  0.])
+    """
+    a = tensor(a)
+    tensor_order = get_order(order, a.order)
+    gpu = a.op.gpu if gpu is None else gpu
+    op = TensorZerosLike(dtype=dtype, gpu=gpu, sparse=a.issparse(), order=order)
+    return op(a, order=tensor_order)

maxframe/tensor/extensions/__init__.py

@@ -0,0 +1,33 @@
+# 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 ...core import CachedAccessor
+from .accessor import TensorMaxFrameAccessor
+from .apply_chunk import apply_chunk
+from .rebalance import rebalance
+
+
+def _install():
+    from ..core import TENSOR_TYPE
+
+    TensorMaxFrameAccessor._register("apply_chunk", apply_chunk)
+    TensorMaxFrameAccessor._register("rebalance", rebalance)
+
+    if TensorMaxFrameAccessor._api_count:
+        for t in TENSOR_TYPE:
+            t.mf = CachedAccessor("mf", TensorMaxFrameAccessor)
+
+
+_install()
+del _install

maxframe/tensor/extensions/accessor.py

@@ -0,0 +1,25 @@
+# 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 TYPE_CHECKING
+
+from ...core import BaseMaxFrameAccessor
+
+if TYPE_CHECKING:
+    from ..core import Tensor
+
+
+class TensorMaxFrameAccessor(BaseMaxFrameAccessor):
+    _obj: "Tensor"
+    _api_count: int = 0

maxframe/tensor/extensions/apply_chunk.py

@@ -0,0 +1,137 @@
+# 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 ENTITY_TYPE
+from ...serialization.serializables import (
+    BoolField,
+    DictField,
+    FunctionField,
+    TupleField,
+)
+from ...udf import BuiltinFunction
+from ...utils import find_objects, quiet_stdio, replace_objects
+from ..core import TensorOrder
+from ..operators import TensorOperator, TensorOperatorMixin
+
+
+class TensorApplyChunk(TensorOperator, TensorOperatorMixin):
+    _op_type_ = opcodes.APPLY_CHUNK
+
+    func = FunctionField("func")
+    elementwise = BoolField("elementwise")
+    args = TupleField("args", default=None)
+    kwargs = DictField("kwargs", default=None)
+    with_chunk_index = BoolField("with_chunk_index", default=False)
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorApplyChunk", inputs):
+        super()._set_inputs(op, inputs)
+        old_inputs = find_objects(op.args, ENTITY_TYPE) + find_objects(
+            op.kwargs, ENTITY_TYPE
+        )
+        mapping = {o: n for o, n in zip(old_inputs, op._inputs[1:])}
+        op.args = replace_objects(op.args, mapping)
+        op.kwargs = replace_objects(op.kwargs, mapping)
+
+    def has_custom_code(self) -> bool:
+        return not isinstance(self.func, BuiltinFunction)
+
+    def __call__(self, t, dtype=None, shape=None, order=None):
+        if dtype is None:
+            try:
+                kwargs = self.kwargs or dict()
+                if self.with_chunk_index:
+                    kwargs["chunk_index"] = (0,) * t.ndim
+                with np.errstate(all="ignore"), quiet_stdio():
+                    mock_result = self.func(
+                        np.random.rand(2, 2).astype(t.dtype),
+                        *(self.args or ()),
+                        **kwargs
+                    )
+            except:
+                raise TypeError("Cannot estimate output type of apply_chunk call")
+            dtype = mock_result.dtype
+            order = (
+                TensorOrder.C_ORDER
+                if mock_result.flags["C_CONTIGUOUS"]
+                else TensorOrder.F_ORDER
+            )
+
+        if shape is not None:
+            new_shape = shape
+        else:
+            new_shape = t.shape if self.elementwise else (np.nan,) * t.ndim
+        inputs = (
+            [t]
+            + find_objects(self.args, ENTITY_TYPE)
+            + find_objects(self.kwargs, ENTITY_TYPE)
+        )
+        return self.new_tensor(inputs, dtype=dtype, shape=new_shape, order=order)
+
+
+def apply_chunk(t, func, args=(), **kwargs):
+    """
+    Apply function to each chunk.
+
+    Parameters
+    ----------
+    func : function
+        Function to apply to each chunk.
+    args : tuple
+        Positional arguments to pass to func in addition to the array.
+    **kwargs
+        Additional keyword arguments to pass as keywords arguments to func.
+
+    Returns
+    -------
+    Tensor
+        Result of applying ``func`` to each chunk of the Tensor.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> a = mt.array([[4, 9]] * 3)
+    >>> a.execute()
+    array([[4, 9],
+           [4, 9],
+           [4, 9]])
+
+    Output dtype will be auto inferred.
+
+    >>> a.mf.apply_chunk(lambda c: c * 0.5).execute()
+    array([[2. , 4.5],
+           [2. , 4.5],
+           [2. , 4.5]])
+
+    You can specify ``dtype`` by yourself if auto infer failed.
+    """
+    elementwise = kwargs.pop("elementwise", None)
+    dtype = np.dtype(kwargs.pop("dtype")) if "dtype" in kwargs else None
+    shape = kwargs.pop("shape", None)
+    order = kwargs.pop("order", None)
+    sparse = kwargs.pop("sparse", t.issparse())
+    with_chunk_index = kwargs.pop("with_chunk_index", False)
+
+    op = TensorApplyChunk(
+        func=func,
+        args=args,
+        kwargs=kwargs,
+        elementwise=elementwise,
+        with_chunk_index=with_chunk_index,
+        sparse=sparse,
+    )
+    return op(t, dtype=dtype, shape=shape, order=order)

maxframe/tensor/extensions/rebalance.py

@@ -0,0 +1,65 @@
+# 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 opcodes
+from ...serialization.serializables import Float64Field, Int64Field
+from ..operators import TensorOperator, TensorOperatorMixin
+
+
+class RebalanceMixin:
+    def __call__(self, in_obj):
+        self._output_types = in_obj.op.output_types
+        return self.new_tileable([in_obj], kws=[in_obj.params])
+
+
+class TensorRebalance(RebalanceMixin, TensorOperatorMixin, TensorOperator):
+    _op_type_ = opcodes.REBALANCE
+
+    axis = Int64Field("axis", default=None)
+    factor = Float64Field("factor", default=None)
+    num_partitions = Int64Field("num_partitions", default=None)
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+
+def rebalance(tensor, axis=0, factor=None, num_partitions=None):
+    """
+    Make data more balanced across entire cluster.
+
+    Parameters
+    ----------
+    axis : int
+        The axis to rebalance.
+    factor : float
+        Specified so that number of chunks after balance is
+        total number of input chunks * factor.
+    num_partitions : int
+        Specified so the number of chunks are at most
+        num_partitions.
+
+    Returns
+    -------
+    Tensor
+        Result of tensor after rebalanced.
+    """
+    if num_partitions is None and factor is None:
+        raise ValueError("Need to specify num_partitions or factor")
+    if num_partitions is not None and factor is not None:
+        raise ValueError(
+            "num_partitions and factor cannot be specified at the same time"
+        )
+
+    op = TensorRebalance(axis=axis, factor=factor, num_partitions=num_partitions)
+    return op(tensor)

maxframe/tensor/fetch/__init__.py

@@ -0,0 +1,15 @@
+# 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 .core import TensorFetch, TensorFetchShuffle

maxframe/tensor/fetch/core.py

@@ -0,0 +1,54 @@
+# 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 ...core import OutputType, register_fetch_class
+from ...core.operator import Fetch, FetchMixin, FetchShuffle
+from ...serialization.serializables import DataTypeField
+from ..operators import TensorOperatorMixin
+
+
+class TensorFetchMixin(TensorOperatorMixin, FetchMixin):
+    __slots__ = ()
+    _output_type_ = OutputType.tensor
+
+
+class TensorFetch(TensorFetchMixin, Fetch):
+    dtype = DataTypeField("dtype")
+
+    def __init__(self, **kw):
+        kw.pop("output_types", None)
+        kw.pop("_output_types", None)
+        super().__init__(**kw)
+
+    def _new_tileables(self, inputs, kws=None, **kw):
+        if "_key" in kw and self.source_key is None:
+            self.source_key = kw["_key"]
+        return super()._new_tileables(inputs, kws=kws, **kw)
+
+
+class TensorFetchShuffle(TensorFetchMixin, FetchShuffle):
+    _dtype = DataTypeField("dtype")
+
+    def __init__(self, **kw):
+        kw.pop("output_types", None)
+        kw.pop("_output_types", None)
+        super().__init__(**kw)
+
+    @property
+    def dtype(self):
+        return getattr(self, "_dtype", None)
+
+
+register_fetch_class(OutputType.tensor, TensorFetch, TensorFetchShuffle)
+register_fetch_class(OutputType.scalar, TensorFetch, TensorFetchShuffle)

maxframe/tensor/fft/__init__.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 .fft import TensorFFT, fft
+from .fft2 import TensorFFT2, fft2
+from .fftfreq import TensorFFTFreq, fftfreq
+from .fftn import TensorFFTN, fftn
+from .fftshift import TensorFFTShift, fftshift
+from .hfft import TensorHFFT, hfft
+from .ifft import TensorIFFT, ifft
+from .ifft2 import TensorIFFT2, ifft2
+from .ifftn import TensorIFFTN, ifftn
+from .ifftshift import TensorIFFTShift, ifftshift
+from .ihfft import TensorIHFFT, ihfft
+from .irfft import TensorIRFFT, irfft
+from .irfft2 import TensorIRFFT2, irfft2
+from .irfftn import TensorIRFFTN, irfftn
+from .rfft import TensorRFFT, rfft
+from .rfft2 import TensorRFFT2, rfft2
+from .rfftfreq import TensorRFFTFreq, rfftfreq
+from .rfftn import TensorRFFTN, rfftn