maxframe 0.1.0b5__cp311-cp311-macosx_11_0_arm64.whl

maxframe/tensor/base/tests/test_base.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 1999-2024 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
+import pytest
+
+from ...datasource import arange, ones, tensor
+from .. import broadcast_to, ravel, transpose, where
+
+
+def test_dir():
+    a = tensor([0, 1, 2], chunk_size=2)
+    tensor_dir = dir(a)
+    for attr in dir(a.data):
+        assert attr in tensor_dir
+
+
+def test_astype():
+    arr = ones((10, 20, 30), chunk_size=3)
+
+    arr2 = arr.astype(np.int32)
+
+    assert arr2.shape == (10, 20, 30)
+    assert np.issubdtype(arr2.dtype, np.int32) is True
+    assert arr2.op.casting == "unsafe"
+
+    with pytest.raises(TypeError):
+        arr.astype(np.int32, casting="safe")
+
+    arr3 = arr.astype(arr.dtype, order="F")
+    assert arr3.flags["F_CONTIGUOUS"] is True
+    assert arr3.flags["C_CONTIGUOUS"] is False
+
+
+def test_transpose():
+    arr = ones((10, 20, 30), chunk_size=[4, 3, 5])
+
+    arr2 = transpose(arr)
+
+    assert arr2.shape == (30, 20, 10)
+
+    with pytest.raises(ValueError):
+        transpose(arr, axes=(1, 0))
+
+    arr3 = transpose(arr, (-2, 2, 0))
+
+    assert arr3.shape == (20, 30, 10)
+
+    arr4 = arr.transpose(-2, 2, 0)
+
+    assert arr4.shape == (20, 30, 10)
+
+    arr5 = arr.T
+
+    assert arr5.shape == (30, 20, 10)
+
+
+def test_broadcast_to():
+    arr = ones((10, 5), chunk_size=2)
+    arr2 = broadcast_to(arr, (20, 10, 5))
+
+    assert arr2.shape == (20, 10, 5)
+
+    arr = ones((10, 5, 1), chunk_size=2)
+    arr3 = broadcast_to(arr, (5, 10, 5, 6))
+
+    assert arr3.shape == (5, 10, 5, 6)
+
+    arr = ones((10, 1), chunk_size=2)
+    arr4 = broadcast_to(arr, (20, 10, 5))
+
+    assert arr4.shape == (20, 10, 5)
+
+    with pytest.raises(ValueError):
+        broadcast_to(arr, (10,))
+
+    with pytest.raises(ValueError):
+        broadcast_to(arr, (5, 1))
+
+    arr = ones((4, 5), chunk_size=2)
+    with pytest.raises((ValueError)):
+        broadcast_to(arr[arr < 2], (3, 20))
+
+
+def test_where():
+    cond = tensor([[True, False], [False, True]], chunk_size=1)
+    x = tensor([1, 2], chunk_size=1)
+
+    with pytest.raises(ValueError):
+        where(cond, x)
+
+    x = arange(9.0).reshape(3, 3)
+    y = where(x < 5, x, -1)
+
+    assert y.dtype == np.float64
+
+
+def test_ravel():
+    arr = ones((10, 5), chunk_size=2)
+    flat_arr = ravel(arr)
+    assert flat_arr.shape == (50,)

maxframe/tensor/base/transpose.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 1999-2024 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 FieldTypes, KeyField, ListField
+from ..core import TensorOrder
+from ..datasource import tensor as astensor
+from ..operators import TensorHasInput, TensorOperatorMixin
+from ..utils import reverse_order
+
+
+def _reorder(x, axes):
+    if x is None:
+        return
+    return type(x)(x[ax] for ax in axes)
+
+
+class TensorTranspose(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.TRANSPOSE
+
+    _input = KeyField("input")
+    axes = ListField("axes", FieldTypes.int32, default=None)
+
+    def __init__(self, axes=None, **kw):
+        # transpose will create a view
+        super().__init__(axes=axes, create_view=True, **kw)
+
+    def __call__(self, a):
+        shape = tuple(
+            s if np.isnan(s) else int(s) for s in _reorder(a.shape, self.axes)
+        )
+        if self.axes == list(reversed(range(a.ndim))):
+            # order reversed
+            tensor_order = reverse_order(a.order)
+        else:
+            tensor_order = TensorOrder.C_ORDER
+        return self.new_tensor([a], shape, order=tensor_order)
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        self._input = self._inputs[0]
+
+    def on_output_modify(self, new_output):
+        op = self.copy().reset_key()
+        return op(new_output)
+
+    def on_input_modify(self, new_input):
+        op = self.copy().reset_key()
+        return op(new_input)
+
+
+def transpose(a, axes=None):
+    """
+    Permute the dimensions of a tensor.
+
+    Parameters
+    ----------
+    a : array_like
+        Input tensor.
+    axes : list of ints, optional
+        By default, reverse the dimensions, otherwise permute the axes
+        according to the values given.
+
+    Returns
+    -------
+    p : Tensor
+        `a` with its axes permuted.  A view is returned whenever
+        possible.
+
+    See Also
+    --------
+    moveaxis
+    argsort
+
+    Notes
+    -----
+    Use `transpose(a, argsort(axes))` to invert the transposition of tensors
+    when using the `axes` keyword argument.
+
+    Transposing a 1-D array returns an unchanged view of the original tensor.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> x = mt.arange(4).reshape((2,2))
+    >>> x.execute()
+    array([[0, 1],
+           [2, 3]])
+
+    >>> mt.transpose(x).execute()
+    array([[0, 2],
+           [1, 3]])
+
+    >>> x = mt.ones((1, 2, 3))
+    >>> mt.transpose(x, (1, 0, 2)).shape
+    (2, 1, 3)
+
+    """
+    a = astensor(a)
+    if axes:
+        if len(axes) != a.ndim:
+            raise ValueError("axes don't match tensor")
+
+    if not axes:
+        axes = list(range(a.ndim))[::-1]
+    else:
+        axes = list(axes)
+    op = TensorTranspose(axes, dtype=a.dtype, sparse=a.issparse())
+    return op(a)

maxframe/tensor/base/unique.py

@@ -0,0 +1,205 @@
+# Copyright 1999-2024 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 as OperandDef
+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_ = OperandDef.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
+
+    def _gen_kws(self, input_obj, chunk=False, chunk_index=None):
+        kws = []
+
+        # unique tensor
+        shape = list(input_obj.shape)
+        shape[self.axis] = np.nan
+        kw = {"shape": tuple(shape), "dtype": input_obj.dtype, "gpu": input_obj.op.gpu}
+        if chunk:
+            idx = [0] * len(shape)
+            idx[self.axis] = chunk_index or 0
+            kw["index"] = tuple(idx)
+        kws.append(kw)
+
+        # unique indices tensor
+        if self.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 self.return_inverse:
+            kw = {
+                "shape": (input_obj.shape[self.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 self.return_counts:
+            kw = {
+                "shape": (np.nan,),
+                "dtype": np.dtype(np.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/base/where.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 1999-2024 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 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)
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        self.condition = self._inputs[0]
+        self.x = self._inputs[1]
+        self.y = self._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)
+        return op(condition, x, y, shape=shape)