maxframe 1.3.1__cp39-cp39-win32.whl → 2.0.0b1__cp39-cp39-win32.whl

maxframe/tensor/datasource/linspace.py

@@ -0,0 +1,140 @@
+# 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 ExecutableTuple
+from ...serialization.serializables import AnyField, BoolField, Int64Field
+from .core import TensorNoInput
+
+
+class TensorLinspace(TensorNoInput):
+    _op_type_ = opcodes.TENSOR_LINSPACE
+
+    start = AnyField("start")
+    stop = AnyField("stop")
+    num = Int64Field("num")
+    endpoint = BoolField("endpoint")
+
+    def __init__(self, dtype=None, **kw):
+        dtype = np.dtype(np.linspace(0, 1, 1).dtype if dtype is None else dtype)
+        super().__init__(dtype=dtype, **kw)
+
+
+def linspace(
+    start,
+    stop,
+    num=50,
+    endpoint=True,
+    retstep=False,
+    dtype=None,
+    gpu=None,
+    chunk_size=None,
+):
+    """
+    Return evenly spaced numbers over a specified interval.
+
+    Returns `num` evenly spaced samples, calculated over the
+    interval [`start`, `stop`].
+
+    The endpoint of the interval can optionally be excluded.
+
+    Parameters
+    ----------
+    start : scalar
+        The starting value of the sequence.
+    stop : scalar
+        The end value of the sequence, unless `endpoint` is set to False.
+        In that case, the sequence consists of all but the last of ``num + 1``
+        evenly spaced samples, so that `stop` is excluded.  Note that the step
+        size changes when `endpoint` is False.
+    num : int, optional
+        Number of samples to generate. Default is 50. Must be non-negative.
+    endpoint : bool, optional
+        If True, `stop` is the last sample. Otherwise, it is not included.
+        Default is True.
+    retstep : bool, optional
+        If True, return (`samples`, `step`), where `step` is the spacing
+        between samples.
+    dtype : dtype, optional
+        The type of the output tensor.  If `dtype` is not given, infer the data
+        type from the other input arguments.
+    gpu : bool, optional
+        Allocate the tensor on GPU if True, False as default
+    chunk_size : int or tuple of int or tuple of ints, optional
+        Desired chunk size on each dimension
+
+    Returns
+    -------
+    samples : Tensor
+        There are `num` equally spaced samples in the closed interval
+        ``[start, stop]`` or the half-open interval ``[start, stop)``
+        (depending on whether `endpoint` is True or False).
+    step : float, optional
+        Only returned if `retstep` is True
+
+        Size of spacing between samples.
+
+
+    See Also
+    --------
+    arange : Similar to `linspace`, but uses a step size (instead of the
+             number of samples).
+    logspace : Samples uniformly distributed in log space.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> mt.linspace(2.0, 3.0, num=5).execute()
+    array([ 2.  ,  2.25,  2.5 ,  2.75,  3.  ])
+    >>> mt.linspace(2.0, 3.0, num=5, endpoint=False).execute()
+    array([ 2. ,  2.2,  2.4,  2.6,  2.8])
+    >>> mt.linspace(2.0, 3.0, num=5, retstep=True).execute()
+    (array([ 2.  ,  2.25,  2.5 ,  2.75,  3.  ]), 0.25)
+
+    Graphical illustration:
+
+    >>> import matplotlib.pyplot as plt
+    >>> N = 8
+    >>> y = mt.zeros(N)
+    >>> x1 = mt.linspace(0, 10, N, endpoint=True)
+    >>> x2 = mt.linspace(0, 10, N, endpoint=False)
+    >>> plt.plot(x1.execute(), y.execute(), 'o')
+    [<matplotlib.lines.Line2D object at 0x...>]
+    >>> plt.plot(x2.execute(), y.execute() + 0.5, 'o')
+    [<matplotlib.lines.Line2D object at 0x...>]
+    >>> plt.ylim([-0.5, 1])
+    (-0.5, 1)
+    >>> plt.show()
+
+    """
+    num = int(num)
+
+    op = TensorLinspace(
+        start=start, stop=stop, num=num, endpoint=endpoint, dtype=dtype, gpu=gpu
+    )
+    shape = (num,)
+    ret = op(shape, chunk_size=chunk_size)
+
+    if not retstep:
+        return ret
+
+    if num > 1:
+        step = float(stop - start) / (num if not endpoint else num - 1)
+    else:
+        step = np.nan
+
+    return ExecutableTuple([ret, step])

maxframe/tensor/datasource/meshgrid.py

@@ -0,0 +1,135 @@
+# 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 .array import tensor
+
+
+def meshgrid(*xi, **kwargs):
+    """
+    Return coordinate matrices from coordinate vectors.
+
+    Make N-D coordinate arrays for vectorized evaluations of
+    N-D scalar/vector fields over N-D grids, given
+    one-dimensional coordinate tensors x1, x2,..., xn.
+
+    Parameters
+    ----------
+    x1, x2,..., xn : array_like
+        1-D arrays representing the coordinates of a grid.
+    indexing : {'xy', 'ij'}, optional
+        Cartesian ('xy', default) or matrix ('ij') indexing of output.
+        See Notes for more details.
+    sparse : bool, optional
+        If True a sparse grid is returned in order to conserve memory.
+        Default is False.
+
+    Returns
+    -------
+    X1, X2,..., XN : Tensor
+        For vectors `x1`, `x2`,..., 'xn' with lengths ``Ni=len(xi)`` ,
+        return ``(N1, N2, N3,...Nn)`` shaped tensors if indexing='ij'
+        or ``(N2, N1, N3,...Nn)`` shaped tensors if indexing='xy'
+        with the elements of `xi` repeated to fill the matrix along
+        the first dimension for `x1`, the second for `x2` and so on.
+
+    Notes
+    -----
+    This function supports both indexing conventions through the indexing
+    keyword argument.  Giving the string 'ij' returns a meshgrid with
+    matrix indexing, while 'xy' returns a meshgrid with Cartesian indexing.
+    In the 2-D case with inputs of length M and N, the outputs are of shape
+    (N, M) for 'xy' indexing and (M, N) for 'ij' indexing.  In the 3-D case
+    with inputs of length M, N and P, outputs are of shape (N, M, P) for
+    'xy' indexing and (M, N, P) for 'ij' indexing.  The difference is
+    illustrated by the following code snippet::
+
+        xv, yv = mt.meshgrid(x, y, sparse=False, indexing='ij')
+        for i in range(nx):
+            for j in range(ny):
+                # treat xv[i,j], yv[i,j]
+
+        xv, yv = mt.meshgrid(x, y, sparse=False, indexing='xy')
+        for i in range(nx):
+            for j in range(ny):
+                # treat xv[j,i], yv[j,i]
+
+    In the 1-D and 0-D case, the indexing and sparse keywords have no effect.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> nx, ny = (3, 2)
+    >>> x = mt.linspace(0, 1, nx)
+    >>> y = mt.linspace(0, 1, ny)
+    >>> xv, yv = mt.meshgrid(x, y)
+    >>> xv.execute()
+    array([[ 0. ,  0.5,  1. ],
+           [ 0. ,  0.5,  1. ]])
+    >>> yv.execute()
+    array([[ 0.,  0.,  0.],
+           [ 1.,  1.,  1.]])
+    >>> xv, yv = mt.meshgrid(x, y, sparse=True)  # make sparse output arrays
+    >>> xv.execute()
+    array([[ 0. ,  0.5,  1. ]])
+    >>> yv.execute()
+    array([[ 0.],
+           [ 1.]])
+
+    `meshgrid` is very useful to evaluate functions on a grid.
+
+    >>> import matplotlib.pyplot as plt
+    >>> x = mt.arange(-5, 5, 0.1)
+    >>> y = mt.arange(-5, 5, 0.1)
+    >>> xx, yy = mt.meshgrid(x, y, sparse=True)
+    >>> z = mt.sin(xx**2 + yy**2) / (xx**2 + yy**2)
+    >>> h = plt.contourf(x,y,z)
+
+    """
+    from ..misc import broadcast_to
+
+    indexing = kwargs.pop("indexing", "xy")
+    sparse = kwargs.pop("sparse", False)
+
+    if kwargs:
+        raise TypeError(
+            f"meshgrid() got an unexpected keyword argument '{list(kwargs)[0]}'"
+        )
+    if indexing not in ("xy", "ij"):
+        raise ValueError("Valid values for `indexing` are 'xy' and 'ij'.")
+
+    xi = [tensor(x) for x in xi]
+    xi = [a.ravel() for a in xi]
+    shape = [x.size for x in xi]
+
+    if indexing == "xy" and len(xi) > 1:
+        xi[0], xi[1] = xi[1], xi[0]
+        shape[0], shape[1] = shape[1], shape[0]
+
+    grid = []
+    for i, x in enumerate(xi):
+        slc = [None] * len(shape)
+        slc[i] = slice(None)
+
+        r = x[tuple(slc)]
+
+        if not sparse:
+            r = broadcast_to(r, shape)
+
+        grid.append(r)
+
+    if indexing == "xy" and len(xi) > 1:
+        grid[0], grid[1] = grid[1], grid[0]
+
+    return grid

maxframe/tensor/datasource/ones.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");
@@ -24,6 +22,7 @@ from ...serialization.serializables import (
     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
@@ -33,7 +32,12 @@ class TensorOnes(TensorNoInput):
     _op_type_ = opcodes.TENSOR_ONES
 
     order = StringField("order")
-    shape = TupleField("shape", FieldTypes.int64)
+    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):
@@ -108,6 +112,7 @@ class TensorOnesLike(TensorLike):
     _op_type_ = opcodes.TENSOR_ONES_LIKE
 
     _input = KeyField("input")
+    order = StringField("order")
 
     def __init__(self, dtype=None, sparse=False, **kw):
         dtype = np.dtype(dtype) if dtype is not None else None

maxframe/tensor/datasource/tests/test_datasource.py

@@ -20,7 +20,7 @@ import scipy.sparse as sps
 
 from .... import dataframe as md
 from ....core import enter_mode
-from ... import arange, full, ones, ones_like, tensor, zeros
+from ... import arange, diag, full, ones, ones_like, tensor, zeros
 from ...core import SparseTensor, Tensor
 from .. import array, asarray, ascontiguousarray, asfortranarray, fromdense
 from ..array import CSRMatrixDataSource
@@ -276,3 +276,34 @@ def test_from_dataframe():
     tensor = from_dataframe(mdf)
     assert tensor.shape == (3, 3)
     assert np.float64 == tensor.dtype
+
+
+def test_diag():
+    # test 2-d, shape[0] == shape[1], k == 0
+    v = tensor(np.arange(16).reshape(4, 4), chunk_size=2)
+    t = diag(v)
+    assert t.shape == (4,)
+
+    v = tensor(np.arange(16).reshape(4, 4), chunk_size=(2, 3))
+    t = diag(v)
+    assert t.shape == (4,)
+
+    # test 1-d, k == 0
+    v = tensor(np.arange(3), chunk_size=2)
+    t = diag(v, sparse=True)
+    assert t.shape == (3, 3)
+
+    # test 2-d, shape[0] != shape[1]
+    v = tensor(np.arange(24).reshape(4, 6), chunk_size=2)
+    t = diag(v)
+    assert t.shape == np.diag(np.arange(24).reshape(4, 6)).shape
+
+    v = tensor(np.arange(24).reshape(4, 6), chunk_size=2)
+    t = diag(v, k=1)
+    assert t.shape == np.diag(np.arange(24).reshape(4, 6), k=1).shape
+    t = diag(v, k=2)
+    assert t.shape == np.diag(np.arange(24).reshape(4, 6), k=2).shape
+    t = diag(v, k=-1)
+    assert t.shape == np.diag(np.arange(24).reshape(4, 6), k=-1).shape
+    t = diag(v, k=-2)
+    assert t.shape == np.diag(np.arange(24).reshape(4, 6), k=-2).shape

maxframe/tensor/datasource/tri_array.py

@@ -0,0 +1,107 @@
+# 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 Int32Field
+from ..core import TensorOrder
+from .array import tensor as astensor
+from .core import TensorHasInput
+
+
+class TensorTriArray(TensorHasInput):
+    k = Int32Field("k")
+
+    def __call__(self, m, order=None):
+        order = TensorOrder.C_ORDER if order is None else order
+        return self.new_tensor([m], shape=m.shape, order=order)
+
+
+class TensorTriu(TensorTriArray):
+    _op_type_ = opcodes.TENSOR_TRIU
+
+
+def triu(m, k=0, gpu=None):
+    """
+    Upper triangle of a tensor.
+
+    Return a copy of a matrix with the elements below the `k`-th diagonal
+    zeroed.
+
+    Please refer to the documentation for `tril` for further details.
+
+    See Also
+    --------
+    tril : lower triangle of a tensor
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> mt.triu([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1).execute()
+    array([[ 1,  2,  3],
+           [ 4,  5,  6],
+           [ 0,  8,  9],
+           [ 0,  0, 12]])
+
+    """
+    m = astensor(m)
+    gpu = m.op.gpu if gpu is None else gpu
+    op = TensorTriu(k=k, dtype=m.dtype, sparse=m.issparse(), gpu=gpu)
+    return op(m)
+
+
+class TensorTril(TensorTriArray):
+    _op_type_ = opcodes.TENSOR_TRIL
+
+
+def tril(m, k=0, gpu=None):
+    """
+    Lower triangle of a tensor.
+
+    Return a copy of a tensor with elements above the `k`-th diagonal zeroed.
+
+    Parameters
+    ----------
+    m : array_like, shape (M, N)
+        Input tensor.
+    k : int, optional
+        Diagonal above which to zero elements.  `k = 0` (the default) is the
+        main diagonal, `k < 0` is below it and `k > 0` is above.
+    gpu : bool, optional
+        Allocate the tensor on GPU if True, None as default
+
+    Returns
+    -------
+    tril : Tensor, shape (M, N)
+        Lower triangle of `m`, of same shape and data-type as `m`.
+
+    See Also
+    --------
+    triu : same thing, only for the upper triangle
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+
+    >>> mt.tril([[1,2,3],[4,5,6],[7,8,9],[10,11,12]], -1).execute()
+    array([[ 0,  0,  0],
+           [ 4,  0,  0],
+           [ 7,  8,  0],
+           [10, 11, 12]])
+
+    """
+    m = astensor(m)
+    gpu = m.op.gpu if gpu is None else gpu
+    op = TensorTril(k=k, dtype=m.dtype, sparse=m.issparse(), gpu=gpu)
+    return op(m)

maxframe/tensor/datasource/zeros.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");
@@ -25,6 +23,7 @@ from ...serialization.serializables import (
     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
@@ -34,7 +33,12 @@ class TensorZeros(TensorNoInput):
     _op_type_ = opcodes.TENSOR_ZEROS
 
     order = StringField("order")
-    shape = TupleField("shape", FieldTypes.int64)
+    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):

maxframe/tensor/extensions/__init__.py

@@ -0,0 +1,31 @@
+# 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
+
+
+def _install():
+    from ..core import TENSOR_TYPE
+
+    TensorMaxFrameAccessor._register("apply_chunk", apply_chunk)
+
+    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")
+    kwargs = DictField("kwargs")
+    with_chunk_index = BoolField("with_chunk_index")
+
+    @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/indexing/__init__.py

@@ -17,7 +17,7 @@ from .compress import compress
 from .extract import extract
 from .fill_diagonal import TensorFillDiagonal, fill_diagonal
 from .flatnonzero import flatnonzero
-from .getitem import FancyIndexingConcat, FancyIndexingDistribute, TensorIndex
+from .getitem import TensorIndex
 from .nonzero import TensorNonzero, nonzero
 from .setitem import TensorIndexSetValue
 from .slice import TensorSlice