maxframe 2.3.0__cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

maxframe/tensor/misc/unique.py

@@ -0,0 +1,205 @@
+# 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 BoolField, Int32Field
+from ..core import TensorOrder
+from ..operators import TensorHasInput, TensorOperatorMixin
+from ..utils import validate_axis
+
+
+class TensorUnique(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.UNIQUE
+
+    return_index = BoolField("return_index", default=False)
+    return_inverse = BoolField("return_inverse", default=False)
+    return_counts = BoolField("return_counts", default=False)
+    axis = Int32Field("axis", default=None)
+
+    @property
+    def output_limit(self):
+        return 1 + self.return_index + self.return_inverse + self.return_counts
+
+    @classmethod
+    def _gen_kws(cls, op: "TensorUnique", input_obj, chunk=False, chunk_index=None):
+        kws = []
+
+        # unique tensor
+        shape = list(input_obj.shape)
+        shape[op.axis] = np.nan
+        kw = {"shape": tuple(shape), "dtype": input_obj.dtype, "gpu": input_obj.op.gpu}
+        if chunk:
+            idx = [0] * len(shape)
+            idx[op.axis] = chunk_index or 0
+            kw["index"] = tuple(idx)
+        kws.append(kw)
+
+        # unique indices tensor
+        if op.return_index:
+            kw = {
+                "shape": (np.nan,),
+                "dtype": np.dtype(np.intp),
+                "gpu": input_obj.op.gpu,
+                "type": "indices",
+            }
+            if chunk:
+                kw["index"] = (chunk_index or 0,)
+            kws.append(kw)
+
+        # unique inverse tensor
+        if op.return_inverse:
+            kw = {
+                "shape": (input_obj.shape[op.axis],),
+                "dtype": np.dtype(np.intp),
+                "gpu": input_obj.op.gpu,
+                "type": "inverse",
+            }
+            if chunk:
+                kw["index"] = (chunk_index or 0,)
+            kws.append(kw)
+
+        # unique counts tensor
+        if op.return_counts:
+            kw = {
+                "shape": (np.nan,),
+                "dtype": np.dtype(int),
+                "gpu": input_obj.op.gpu,
+                "type": "counts",
+            }
+            if chunk:
+                kw["index"] = (chunk_index or 0,)
+            kws.append(kw)
+
+        return kws
+
+    def __call__(self, ar):
+        from .atleast_1d import atleast_1d
+
+        ar = atleast_1d(ar)
+        if self.axis is None:
+            if ar.ndim > 1:
+                ar = ar.flatten()
+            self.axis = 0
+        else:
+            self.axis = validate_axis(ar.ndim, self.axis)
+
+        kws = self._gen_kws(self, ar)
+        tensors = self.new_tensors([ar], kws=kws, order=TensorOrder.C_ORDER)
+        if len(tensors) == 1:
+            return tensors[0]
+        return tensors
+
+
+def unique(
+    ar,
+    return_index=False,
+    return_inverse=False,
+    return_counts=False,
+    axis=None,
+):
+    """
+    Find the unique elements of a tensor.
+
+    Returns the sorted unique elements of a tensor. There are three optional
+    outputs in addition to the unique elements:
+
+    * the indices of the input tensor that give the unique values
+    * the indices of the unique tensor that reconstruct the input tensor
+    * the number of times each unique value comes up in the input tensor
+
+    Parameters
+    ----------
+    ar : array_like
+        Input tensor. Unless `axis` is specified, this will be flattened if it
+        is not already 1-D.
+    return_index : bool, optional
+        If True, also return the indices of `ar` (along the specified axis,
+        if provided, or in the flattened tensor) that result in the unique tensor.
+    return_inverse : bool, optional
+        If True, also return the indices of the unique tensor (for the specified
+        axis, if provided) that can be used to reconstruct `ar`.
+    return_counts : bool, optional
+        If True, also return the number of times each unique item appears
+        in `ar`.
+    axis : int or None, optional
+        The axis to operate on. If None, `ar` will be flattened. If an integer,
+        the subarrays indexed by the given axis will be flattened and treated
+        as the elements of a 1-D tensor with the dimension of the given axis,
+        see the notes for more details.  Object tensors or structured tensors
+        that contain objects are not supported if the `axis` kwarg is used. The
+        default is None.
+
+    Returns
+    -------
+    unique : Tensor
+        The sorted unique values.
+    unique_indices : Tensor, optional
+        The indices of the first occurrences of the unique values in the
+        original tensor. Only provided if `return_index` is True.
+    unique_inverse : Tensor, optional
+        The indices to reconstruct the original tensor from the
+        unique tensor. Only provided if `return_inverse` is True.
+    unique_counts : Tensor, optional
+        The number of times each of the unique values comes up in the
+        original tensor. Only provided if `return_counts` is True.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> mt.unique([1, 1, 2, 2, 3, 3]).execute()
+    array([1, 2, 3])
+    >>> a = mt.array([[1, 1], [2, 3]])
+    >>> mt.unique(a).execute()
+    array([1, 2, 3])
+
+    Return the unique rows of a 2D tensor
+
+    >>> a = mt.array([[1, 0, 0], [1, 0, 0], [2, 3, 4]])
+    >>> mt.unique(a, axis=0).execute()
+    array([[1, 0, 0], [2, 3, 4]])
+
+    Return the indices of the original tensor that give the unique values:
+
+    >>> a = mt.array(['a', 'b', 'b', 'c', 'a'])
+    >>> u, indices = mt.unique(a, return_index=True)
+    >>> u.execute()
+    array(['a', 'b', 'c'],
+           dtype='|S1')
+    >>> indices.execute()
+    array([0, 1, 3])
+    >>> a[indices].execute()
+    array(['a', 'b', 'c'],
+           dtype='|S1')
+
+    Reconstruct the input array from the unique values:
+
+    >>> a = mt.array([1, 2, 6, 4, 2, 3, 2])
+    >>> u, indices = mt.unique(a, return_inverse=True)
+    >>> u.execute()
+    array([1, 2, 3, 4, 6])
+    >>> indices.execute()
+    array([0, 1, 4, 3, 1, 2, 1])
+    >>> u[indices].execute()
+    array([1, 2, 6, 4, 2, 3, 2])
+    """
+    op = TensorUnique(
+        return_index=return_index,
+        return_inverse=return_inverse,
+        return_counts=return_counts,
+        axis=axis,
+    )
+    return op(ar)

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/misc/where.py

@@ -0,0 +1,129 @@
+# 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 KeyField
+from ..datasource import tensor as astensor
+from ..operators import TensorOperator, TensorOperatorMixin
+from ..utils import broadcast_shape
+from .broadcast_to import broadcast_to
+
+
+class TensorWhere(TensorOperator, TensorOperatorMixin):
+    _op_type_ = opcodes.WHERE
+
+    condition = KeyField("condition", default=None)
+    x = KeyField("x", default=None)
+    y = KeyField("y", default=None)
+
+    @classmethod
+    def _set_inputs(cls, op: "TensorWhere", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        op.condition = op._inputs[0]
+        op.x = op._inputs[1]
+        op.y = op._inputs[2]
+
+    def __call__(self, condition, x, y, shape=None):
+        shape = shape or broadcast_shape(condition.shape, x.shape, y.shape)
+        return self.new_tensor([condition, x, y], shape)
+
+
+def where(condition, x=None, y=None):
+    """
+    Return elements, either from `x` or `y`, depending on `condition`.
+
+    If only `condition` is given, return ``condition.nonzero()``.
+
+    Parameters
+    ----------
+    condition : array_like, bool
+        When True, yield `x`, otherwise yield `y`.
+    x, y : array_like, optional
+        Values from which to choose. `x`, `y` and `condition` need to be
+        broadcastable to some shape.
+
+    Returns
+    -------
+    out : Tensor or tuple of Tensors
+        If both `x` and `y` are specified, the output tensor contains
+        elements of `x` where `condition` is True, and elements from
+        `y` elsewhere.
+
+        If only `condition` is given, return the tuple
+        ``condition.nonzero()``, the indices where `condition` is True.
+
+    See Also
+    --------
+    nonzero, choose
+
+    Notes
+    -----
+    If `x` and `y` are given and input arrays are 1-D, `where` is
+    equivalent to::
+
+        [xv if c else yv for (c,xv,yv) in zip(condition,x,y)]
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> mt.where([[True, False], [True, True]],
+    ...          [[1, 2], [3, 4]],
+    ...          [[9, 8], [7, 6]]).execute()
+    array([[1, 8],
+           [3, 4]])
+
+    >>> mt.where([[0, 1], [1, 0]]).execute()
+    (array([0, 1]), array([1, 0]))
+
+    >>> x = mt.arange(9.).reshape(3, 3)
+    >>> mt.where( x > 5 ).execute()
+    (array([2, 2, 2]), array([0, 1, 2]))
+    >>> mt.where(x < 5, x, -1).execute()               # Note: broadcasting.
+    array([[ 0.,  1.,  2.],
+           [ 3.,  4., -1.],
+           [-1., -1., -1.]])
+
+    Find the indices of elements of `x` that are in `goodvalues`.
+
+    >>> goodvalues = [3, 4, 7]
+    >>> ix = mt.isin(x, goodvalues)
+    >>> ix.execute()
+    array([[False, False, False],
+           [ True,  True, False],
+           [False,  True, False]])
+    >>> mt.where(ix).execute()
+    (array([1, 1, 2]), array([0, 1, 1]))
+    """
+    if (x is None) != (y is None):
+        raise ValueError("either both or neither of x and y should be given")
+
+    if x is None and y is None:
+        return astensor(condition).nonzero()
+
+    x, y = astensor(x), astensor(y)
+    dtype = np.result_type(x.dtype, y.dtype)
+    shape = broadcast_shape(x.shape, y.shape)
+
+    if np.isscalar(condition):
+        return broadcast_to(x if condition else y, shape).astype(dtype)
+    else:
+        condition = astensor(condition)
+        op = TensorWhere(dtype=dtype, sparse=condition.issparse())
+        return op(condition, x, y, shape=shape)

maxframe/tensor/operators.py

@@ -0,0 +1,83 @@
+# 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
+from ..core.operator import (
+    HasInput,
+    MapReduceOperator,
+    Operator,
+    ShuffleProxy,
+    TileableOperatorMixin,
+)
+from ..serialization.serializables import DataTypeField
+
+
+class TensorOperatorMixin(TileableOperatorMixin):
+    __slots__ = ()
+    _op_module_ = "tensor"
+    _output_type_ = OutputType.tensor
+
+    def new_tensors(
+        self,
+        inputs,
+        shape=None,
+        dtype=None,
+        order=None,
+        chunks=None,
+        nsplits=None,
+        output_limit=None,
+        kws=None,
+        **kw
+    ):
+        return self.new_tileables(
+            inputs,
+            shape=shape,
+            chunks=chunks,
+            nsplits=nsplits,
+            output_limit=output_limit,
+            kws=kws,
+            dtype=dtype,
+            order=order,
+            **kw
+        )
+
+    def new_tensor(self, inputs, shape, dtype=None, order=None, **kw):
+        if getattr(self, "output_limit") != 1:
+            raise TypeError("cannot new tensor with more than 1 outputs")
+        return self.new_tensors(inputs, shape=shape, dtype=dtype, order=order, **kw)[0]
+
+
+class TensorOperator(Operator):
+    _output_type_ = OutputType.tensor
+
+    dtype = DataTypeField("dtype", default=None)
+
+
+class TensorHasInput(HasInput):
+    _output_type_ = OutputType.tensor
+
+    dtype = DataTypeField("dtype", default=None)
+
+
+class TensorShuffleProxy(ShuffleProxy, TensorOperatorMixin):
+    _output_type_ = OutputType.tensor
+
+    dtype = DataTypeField("dtype", default=None)
+
+
+class TensorMapReduceOperator(MapReduceOperator):
+    _output_type_ = OutputType.tensor
+
+    dtype = DataTypeField("dtype", default=None)

maxframe/tensor/random/__init__.py

@@ -0,0 +1,166 @@
+# 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 .beta import TensorRandBeta, beta
+from .binomial import TensorBinomial, binomial
+from .bytes import bytes
+from .chisquare import TensorChisquareDist, chisquare
+from .choice import TensorChoice, choice
+from .core import RandomState, RandomStateField, TensorDistribution, _random_state
+from .dirichlet import TensorDirichlet, dirichlet
+from .exponential import TensorExponential, exponential
+from .f import TensorFDist, f
+from .gamma import TensorRandGamma, gamma
+from .geometric import TensorGeometric, geometric
+from .gumbel import TensorGumbel, gumbel
+from .hypergeometric import TensorHypergeometric, hypergeometric
+from .laplace import TensorLaplaceDist, laplace
+from .logistic import TensorLogistic, logistic
+from .lognormal import TensorLognormal, lognormal
+from .logseries import TensorLogseries, logseries
+from .multinomial import TensorMultinomial, multinomial
+from .multivariate_normal import TensorMultivariateNormal, multivariate_normal
+from .negative_binomial import TensorNegativeBinomial, negative_binomial
+from .noncentral_chisquare import TensorNoncentralChisquare, noncentral_chisquare
+from .noncentral_f import TensorNoncentralF, noncentral_f
+from .normal import TensorNormal, normal
+from .pareto import TensorPareto, pareto
+from .permutation import TensorPermutation, permutation
+from .poisson import TensorPoisson, poisson
+from .power import TensorRandomPower, power
+from .rand import TensorRand, rand
+from .randint import TensorRandint, randint
+from .randn import TensorRandn, randn
+from .random_integers import TensorRandomIntegers, random_integers
+from .random_sample import TensorRandomSample, random_sample
+from .rayleigh import TensorRayleigh, rayleigh
+from .shuffle import shuffle
+from .standard_cauchy import TensorStandardCauchy, standard_cauchy
+from .standard_exponential import TensorStandardExponential, standard_exponential
+from .standard_gamma import TensorStandardGamma, standard_gamma
+from .standard_normal import TensorStandardNormal, standard_normal
+from .standard_t import TensorStandardT, standard_t
+from .triangular import TensorTriangular, triangular
+from .uniform import TensorUniform, uniform
+from .vonmises import TensorVonmises, vonmises
+from .wald import TensorWald, wald
+from .weibull import TensorWeibull, weibull
+from .zipf import TensorZipf, zipf
+
+
+def _install():
+    setattr(RandomState, "rand", rand)
+    setattr(RandomState, "randn", randn)
+    setattr(RandomState, "randint", randint)
+    setattr(RandomState, "random_integers", random_integers)
+    setattr(RandomState, "random_sample", random_sample)
+    setattr(RandomState, "ranf", random_sample)
+    setattr(RandomState, "random", random_sample)
+    setattr(RandomState, "sample", random_sample)
+    setattr(RandomState, "choice", choice)
+    setattr(RandomState, "bytes", bytes)
+    setattr(RandomState, "beta", beta)
+    setattr(RandomState, "binomial", binomial)
+    setattr(RandomState, "chisquare", chisquare)
+    setattr(RandomState, "dirichlet", dirichlet)
+    setattr(RandomState, "exponential", exponential)
+    setattr(RandomState, "f", f)
+    setattr(RandomState, "gamma", gamma)
+    setattr(RandomState, "geometric", geometric)
+    setattr(RandomState, "gumbel", gumbel)
+    setattr(RandomState, "hypergeometric", hypergeometric)
+    setattr(RandomState, "laplace", laplace)
+    setattr(RandomState, "logistic", logistic)
+    setattr(RandomState, "lognormal", lognormal)
+    setattr(RandomState, "logseries", logseries)
+    setattr(RandomState, "multinomial", multinomial)
+    setattr(RandomState, "multivariate_normal", multivariate_normal)
+    setattr(RandomState, "negative_binomial", negative_binomial)
+    setattr(RandomState, "noncentral_chisquare", noncentral_chisquare)
+    setattr(RandomState, "noncentral_f", noncentral_f)
+    setattr(RandomState, "normal", normal)
+    setattr(RandomState, "pareto", pareto)
+    setattr(RandomState, "poisson", poisson)
+    setattr(RandomState, "power", power)
+    setattr(RandomState, "rayleigh", rayleigh)
+    setattr(RandomState, "standard_cauchy", standard_cauchy)
+    setattr(RandomState, "standard_exponential", standard_exponential)
+    setattr(RandomState, "standard_gamma", standard_gamma)
+    setattr(RandomState, "standard_normal", standard_normal)
+    setattr(RandomState, "standard_t", standard_t)
+    setattr(RandomState, "triangular", triangular)
+    setattr(RandomState, "uniform", uniform)
+    setattr(RandomState, "vonmises", vonmises)
+    setattr(RandomState, "wald", wald)
+    setattr(RandomState, "weibull", weibull)
+    setattr(RandomState, "zipf", zipf)
+    setattr(RandomState, "permutation", permutation)
+    setattr(RandomState, "shuffle", shuffle)
+
+
+_install()
+del _install
+
+
+seed = _random_state.seed
+
+rand = _random_state.rand
+randn = _random_state.randn
+randint = _random_state.randint
+random_integers = _random_state.random_integers
+random_sample = _random_state.random_sample
+random = _random_state.random
+ranf = _random_state.ranf
+sample = _random_state.sample
+choice = _random_state.choice
+bytes = _random_state.bytes
+
+permutation = _random_state.permutation
+shuffle = _random_state.shuffle
+
+beta = _random_state.beta
+binomial = _random_state.binomial
+chisquare = _random_state.chisquare
+dirichlet = _random_state.dirichlet
+exponential = _random_state.exponential
+f = _random_state.f
+gamma = _random_state.gamma
+geometric = _random_state.geometric
+gumbel = _random_state.gumbel
+hypergeometric = _random_state.hypergeometric
+laplace = _random_state.laplace
+logistic = _random_state.logistic
+lognormal = _random_state.lognormal
+logseries = _random_state.logseries
+multinomial = _random_state.multinomial
+multivariate_normal = _random_state.multivariate_normal
+negative_binomial = _random_state.negative_binomial
+noncentral_chisquare = _random_state.noncentral_chisquare
+noncentral_f = _random_state.noncentral_f
+normal = _random_state.normal
+pareto = _random_state.pareto
+poisson = _random_state.poisson
+power = _random_state.power
+rayleigh = _random_state.rayleigh
+standard_cauchy = _random_state.standard_cauchy
+standard_exponential = _random_state.standard_exponential
+standard_gamma = _random_state.standard_gamma
+standard_normal = _random_state.standard_normal
+standard_t = _random_state.standard_t
+triangular = _random_state.triangular
+uniform = _random_state.uniform
+vonmises = _random_state.vonmises
+wald = _random_state.wald
+weibull = _random_state.weibull
+zipf = _random_state.zipf