maxframe 0.1.0b5__cp310-cp310-macosx_11_0_arm64.whl

maxframe/tensor/array_utils.py

@@ -0,0 +1,186 @@
+# 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.
+
+from collections import defaultdict
+from contextlib import contextmanager
+
+import numpy as np
+
+from ..lib import sparse
+from ..lib.sparse.core import get_dense_module, issparse
+from ..utils import lazy_import
+
+cp = lazy_import("cupy", rename="cp")
+
+
+def is_array(x):
+    if isinstance(x, np.ndarray):
+        return True
+    elif isinstance(x, (sparse.SparseMatrix, sparse.SparseVector)):
+        return True
+    elif cp:  # pragma: no cover
+        return isinstance(x, cp.ndarray)
+    else:
+        return False
+
+
+def is_cupy(x):
+    if cp and isinstance(x, cp.ndarray):  # pragma: no cover
+        return True
+    else:
+        return False
+
+
+def get_array_module(x, nosparse=False):
+    if issparse(x):
+        if nosparse:
+            return get_dense_module(x)
+        return sparse
+    if cp:
+        return cp.get_array_module(x)
+    return np
+
+
+def array_module(gpu):
+    if gpu:
+        if cp is None:
+            raise ImportError("Execute on GPU requires for `cupy` library")
+        return cp
+
+    return np
+
+
+def _get(x):
+    m = get_array_module(x)
+
+    if m is np:
+        return x
+    if m is sparse:
+        return x if not hasattr(x, "get") else x.get()
+    return x.get()
+
+
+def move_to_device(x, device_id):
+    if hasattr(x, "device") and x.device.id == device_id:
+        return x
+
+    assert device_id >= 0
+
+    if issparse(x) and device_id > 0:
+        raise NotImplementedError
+
+    # for dense array, we currently copy from gpu to memory and then copy back to destination device
+    # to avoid kernel panic
+    with cp.cuda.Device(device_id):
+        return cp.asarray(cp.asnumpy(x))  # remove `cp.asnumpy` call to do directly copy
+
+
+def convert_order(x, order):
+    xp = get_array_module(x)
+    if xp.isfortran(x) != (order == "F"):
+        x = xp.array(x, order=order)
+    return x
+
+
+def _most_nbytes_device(device_nbytes):
+    device_to_nbytes = defaultdict(lambda: 0)
+    for device, nbytes in device_nbytes:
+        device_to_nbytes[device] += nbytes
+    return max(device_to_nbytes, key=lambda i: device_to_nbytes[i])
+
+
+def _is_array_writeable(a):
+    if hasattr(a, "flags") and hasattr(a.flags, "writeable"):
+        return a.flags.writeable
+    # writeable as default
+    return True
+
+
+def as_same_device(inputs, device=None, ret_extra=False, copy_if_not_writeable=False):
+    input_tensors = [
+        i for i in inputs if hasattr(i, "ndim") and i.ndim > 0
+    ]  # filter scalar
+    has_sparse = any(issparse(i) for i in inputs)
+
+    if device is None:
+        try:
+            device = _most_nbytes_device(
+                (i.device.id if hasattr(i, "device") else -1, i.nbytes)
+                for i in input_tensors
+            )
+        except ValueError:
+            device = -1
+
+    if device == -1:
+        outputs = [_get(i) for i in inputs]
+    else:
+        outputs = [move_to_device(i, device) for i in inputs]
+
+    if copy_if_not_writeable:
+        new_outputs = []
+        for out in outputs:
+            if not _is_array_writeable(out):
+                new_outputs.append(out.copy())
+            elif isinstance(out, (sparse.SparseMatrix, sparse.SparseVector)):
+                if (
+                    not _is_array_writeable(out.data)
+                    or not _is_array_writeable(out.indices)
+                    or not _is_array_writeable(out.indptr)
+                ):
+                    new_outputs.append(type(out)(out.spmatrix.copy(), shape=out.shape))
+                else:
+                    new_outputs.append(out)
+            else:
+                new_outputs.append(out)
+        outputs = new_outputs
+
+    if not ret_extra:
+        return outputs
+
+    if has_sparse:
+        m = sparse
+    else:
+        if len(input_tensors) > 0:
+            m = get_array_module(input_tensors[0])
+        else:
+            m = np
+    return outputs, device, m
+
+
+def as_np_array(x):
+    xp = get_array_module(x)
+    return x if xp == np else x.get()
+
+
+def is_sparse_module(xp):
+    return xp is sparse
+
+
+@contextmanager
+def device(device_id):
+    if device_id is None or device_id < 0:
+        yield
+    else:  # pragma: no cover
+        with cp.cuda.Device(device_id) as dev:
+            yield dev
+
+
+def create_array(op):
+    xp = array_module(op.gpu)
+
+    def inner(func, *args, **kwargs):
+        with device(op.device):
+            return getattr(xp, func)(*args, **kwargs)
+
+    return inner

maxframe/tensor/base/__init__.py

@@ -0,0 +1,34 @@
+# 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.
+
+from .astype import TensorAstype
+from .atleast_1d import atleast_1d
+from .broadcast_to import TensorBroadcastTo, broadcast_to
+from .ravel import ravel
+from .transpose import transpose
+from .unique import unique
+from .where import TensorWhere, where
+
+
+def _install():
+    from ..core import Tensor, TensorData
+    from .astype import _astype
+
+    for cls in (Tensor, TensorData):
+        setattr(cls, "astype", _astype)
+        setattr(cls, "ravel", ravel)
+
+
+_install()
+del _install

maxframe/tensor/base/astype.py

@@ -0,0 +1,119 @@
+#!/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, StringField
+from ...utils import get_dtype
+from ..operators import TensorHasInput, TensorOperatorMixin
+from ..utils import get_order
+
+
+class TensorAstype(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.ASTYPE
+
+    _input = KeyField("input")
+    order = StringField("order", default=None)
+    casting = StringField("casting", default=None)
+
+    def __init__(self, dtype=None, sparse=False, **kw):
+        super().__init__(dtype=dtype, sparse=sparse, **kw)
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        self._input = self._inputs[0]
+
+    def __call__(self, tensor, order=None):
+        return self.new_tensor([tensor], tensor.shape, order=order)
+
+
+def _astype(tensor, dtype, order="K", casting="unsafe", copy=True):
+    """
+    Copy of the tensor, cast to a specified type.
+
+    Parameters
+    ----------
+    dtype : str or dtype
+        Typecode or data-type to which the array is cast.
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        Controls what kind of data casting may occur. Defaults to 'unsafe'
+        for backwards compatibility.
+          * 'no' means the data types should not be cast at all.
+          * 'equiv' means only byte-order changes are allowed.
+          * 'safe' means only casts which can preserve values are allowed.
+          * 'same_kind' means only safe casts or casts within a kind,
+            like float64 to float32, are allowed.
+          * 'unsafe' means any data conversions may be done.
+    order : {'C', 'F', 'A', 'K'}, optional
+        Controls the memory layout order of the result.
+        'C' means C order, 'F' means Fortran order, 'A'
+        means 'F' order if all the arrays are Fortran contiguous,
+        'C' order otherwise, and 'K' means as close to the
+        order the array elements appear in memory as possible.
+        Default is 'K'.
+    copy : bool, optional
+        By default, astype always returns a newly allocated array. If this
+        is set to false, and the `dtype`, `order`, and `subok`
+        requirements are satisfied, the input array is returned instead
+        of a copy.
+
+    Returns
+    -------
+    arr_t : Tensor
+        Unless `copy` is False and the other conditions for returning the input
+        array are satisfied (see description for `copy` input parameter), `arr_t`
+        is a new tensor of the same shape as the input array, with dtype, order
+        given by `dtype`, `order`.
+
+    Notes
+    -----
+    astype method returns an error if the string
+    dtype to cast to is not long enough in 'safe' casting mode to hold the max
+    value of integer/float array that is being casted. Previously the casting
+    was allowed even if the result was truncated.
+
+    Raises
+    ------
+    ComplexWarning
+        When casting from complex to float or int. To avoid this,
+        one should use ``a.real.astype(t)``.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> x = mt.array([1, 2, 2.5])
+    >>> x.execute()
+    array([ 1. ,  2. ,  2.5])
+
+    >>> x.astype(int).execute()
+    array([1, 2, 2])
+    """
+    dtype = get_dtype(dtype)
+    tensor_order = get_order(order, tensor.order)
+
+    if tensor.dtype == dtype and tensor.order == tensor_order:
+        return tensor if not copy else tensor.copy(order=order)
+    elif not np.can_cast(tensor.dtype, dtype, casting=casting):
+        raise TypeError(
+            f"Cannot cast array from {tensor.dtype!r} to {dtype!r} "
+            f"according to the rule {casting}"
+        )
+
+    op = TensorAstype(
+        dtype=dtype, order=order, casting=casting, sparse=tensor.issparse()
+    )
+    return op(tensor, order=tensor_order)

maxframe/tensor/base/atleast_1d.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Copyright 1999-2021 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 ...core import ExecutableTuple
+from ..datasource import tensor as astensor
+
+
+def atleast_1d(*tensors):
+    """
+    Convert inputs to tensors with at least one dimension.
+
+    Scalar inputs are converted to 1-dimensional tensors, whilst
+    higher-dimensional inputs are preserved.
+
+    Parameters
+    ----------
+    tensors1, tensors2, ... : array_like
+        One or more input tensors.
+
+    Returns
+    -------
+    ret : Tensor
+        An tensor, or list of tensors, each with ``a.ndim >= 1``.
+        Copies are made only if necessary.
+
+    See Also
+    --------
+    atleast_2d, atleast_3d
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> mt.atleast_1d(1.0).execute()
+    array([ 1.])
+
+    >>> x = mt.arange(9.0).reshape(3,3)
+    >>> mt.atleast_1d(x).execute()
+    array([[ 0.,  1.,  2.],
+           [ 3.,  4.,  5.],
+           [ 6.,  7.,  8.]])
+    >>> mt.atleast_1d(x) is x
+    True
+
+    >>> mt.atleast_1d(1, [3, 4]).execute()
+    [array([1]), array([3, 4])]
+
+    """
+    new_tensors = []
+    for x in tensors:
+        x = astensor(x)
+        if x.ndim == 0:
+            x = x[np.newaxis]
+
+        new_tensors.append(x)
+
+    if len(new_tensors) == 1:
+        return new_tensors[0]
+    return ExecutableTuple(new_tensors)

maxframe/tensor/base/broadcast_to.py

@@ -0,0 +1,89 @@
+# 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, TupleField
+from ..datasource import tensor as astensor
+from ..operators import TensorHasInput, TensorOperatorMixin
+
+
+class TensorBroadcastTo(TensorHasInput, TensorOperatorMixin):
+    _op_type_ = opcodes.BROADCAST_TO
+
+    _input = KeyField("input")
+    shape = TupleField("shape", default=None)
+
+    def __call__(self, tensor, shape):
+        return self.new_tensor([tensor], shape)
+
+
+def broadcast_to(tensor, shape):
+    """Broadcast a tensor to a new shape.
+
+    Parameters
+    ----------
+    tensor : array_like
+        The tensor to broadcast.
+    shape : tuple
+        The shape of the desired array.
+
+    Returns
+    -------
+    broadcast : Tensor
+
+    Raises
+    ------
+    ValueError
+        If the tensor is not compatible with the new shape according to MaxFrame's
+        broadcasting rules.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> x = mt.array([1, 2, 3])
+    >>> mt.broadcast_to(x, (3, 3)).execute()
+    array([[1, 2, 3],
+           [1, 2, 3],
+           [1, 2, 3]])
+    """
+    from ..core import Tensor
+
+    tensor = tensor if isinstance(tensor, Tensor) else astensor(tensor)
+    shape = tuple(shape) if isinstance(shape, (list, tuple)) else (shape,)
+
+    if any(np.isnan(s) for s in tensor.shape):
+        raise ValueError(
+            "input tensor has unknown shape, need to call `.execute()` first"
+        )
+
+    if tensor.shape == shape:
+        return tensor
+
+    new_ndim = len(shape) - tensor.ndim
+    if new_ndim < 0:
+        raise ValueError(
+            "input operator has more dimensions than allowed by the axis remapping"
+        )
+    if any(o != n for o, n in zip(tensor.shape, shape[new_ndim:]) if o != 1):
+        raise ValueError(
+            "operators could not be broadcast together "
+            f"with remapped shapes [original->remapped]: {tensor.shape} "
+            f"and requested shape {shape}"
+        )
+
+    op = TensorBroadcastTo(shape, dtype=tensor.dtype, sparse=tensor.issparse())
+    return op(tensor, shape)

maxframe/tensor/base/ravel.py

@@ -0,0 +1,92 @@
+#!/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 ..datasource import tensor as astensor
+
+
+def ravel(a, order="C"):
+    """Return a contiguous flattened tensor.
+
+    A 1-D tensor, containing the elements of the input, is returned.  A copy is
+    made only if needed.
+
+    Parameters
+    ----------
+    a : array_like
+        Input tensor.  The elements in `a` are packed as a 1-D tensor.
+    order : {'C','F', 'A', 'K'}, optional
+
+        The elements of `a` are read using this index order. 'C' means
+        to index the elements in row-major, C-style order,
+        with the last axis index changing fastest, back to the first
+        axis index changing slowest.  'F' means to index the elements
+        in column-major, Fortran-style order, with the
+        first index changing fastest, and the last index changing
+        slowest. Note that the 'C' and 'F' options take no account of
+        the memory layout of the underlying array, and only refer to
+        the order of axis indexing.  'A' means to read the elements in
+        Fortran-like index order if `a` is Fortran *contiguous* in
+        memory, C-like order otherwise.  'K' means to read the
+        elements in the order they occur in memory, except for
+        reversing the data when strides are negative.  By default, 'C'
+        index order is used.
+
+    Returns
+    -------
+    y : array_like
+        If `a` is a matrix, y is a 1-D tensor, otherwise y is a tensor of
+        the same subtype as `a`. The shape of the returned array is
+        ``(a.size,)``. Matrices are special cased for backward
+        compatibility.
+
+    See Also
+    --------
+    Tensor.flat : 1-D iterator over an array.
+    Tensor.flatten : 1-D array copy of the elements of an array
+                      in row-major order.
+    Tensor.reshape : Change the shape of an array without changing its data.
+
+    Examples
+    --------
+    It is equivalent to ``reshape(-1)``.
+
+    >>> import maxframe.tensor as mt
+
+    >>> x = mt.array([[1, 2, 3], [4, 5, 6]])
+    >>> print(mt.ravel(x).execute())
+    [1 2 3 4 5 6]
+
+    >>> print(x.reshape(-1).execute())
+    [1 2 3 4 5 6]
+
+    >>> print(mt.ravel(x.T).execute())
+    [1 4 2 5 3 6]
+
+    >>> a = mt.arange(12).reshape(2,3,2).swapaxes(1,2); a.execute()
+    array([[[ 0,  2,  4],
+            [ 1,  3,  5]],
+           [[ 6,  8, 10],
+            [ 7,  9, 11]]])
+    >>> a.ravel().execute()
+    array([ 0,  2,  4,  1,  3,  5,  6,  8, 10,  7,  9, 11])
+
+    """
+    a = astensor(a)
+    if a.ndim == 0:
+        return a[np.newaxis]
+    return a.reshape(-1, order=order)

maxframe/tensor/base/tests/__init__.py

@@ -0,0 +1,13 @@
+# 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.