maxframe 2.0.0b2__cp310-cp310-win_amd64.whl → 2.2.0__cp310-cp310-win_amd64.whl

maxframe/codegen/spe/tensor/fft.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 typing import List, Union
+
+from ....tensor.fft import TensorFFTFreq, TensorRFFTFreq
+from ....tensor.fft.core import TensorFFTMixin, TensorFFTNMixin, TensorFFTShiftMixin
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+
+
+@register_op_adapter(TensorFFTMixin)
+class TensorFFTAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: TensorFFTMixin, context: SPECodeContext) -> List[str]:
+        context.register_import("numpy", "np")
+        func_name = type(op).__name__.lower()[6:]
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        kw_str = self.generate_call_args_with_attributes(
+            op, context, kw_keys=["n", "axis", "norm"], skip_none=True
+        )
+        return [f"{res_var_name} = np.fft.{func_name}({input_var_name}, {kw_str})"]
+
+
+@register_op_adapter(TensorFFTNMixin)
+class TensorFFTNAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: TensorFFTNMixin, context: SPECodeContext) -> List[str]:
+        context.register_import("numpy", "np")
+        func_name = type(op).__name__.lower()[6:]
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        kw_str = self.generate_call_args_with_attributes(
+            op, context, "shape", kw_keys=["axes", "norm"], skip_none=True
+        )
+        return [f"{res_var_name} = np.fft.{func_name}({input_var_name}, {kw_str})"]
+
+
+@register_op_adapter(TensorFFTShiftMixin)
+class TensorFFTShiftAdapter(SPEOperatorAdapter):
+    def generate_code(
+        self, op: TensorFFTShiftMixin, context: SPECodeContext
+    ) -> List[str]:
+        context.register_import("numpy", "np")
+        func_name = type(op).__name__.lower()[6:]
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        kw_str = self.generate_call_args_with_attributes(
+            op, context, kw_keys=["axes"], skip_none=True
+        )
+        return [f"{res_var_name} = np.fft.{func_name}({input_var_name}, {kw_str})"]
+
+
+@register_op_adapter([TensorFFTFreq, TensorRFFTFreq])
+class TensorFFTFreqAdapter(SPEOperatorAdapter):
+    def generate_code(
+        self, op: Union[TensorFFTFreq, TensorRFFTFreq], context: SPECodeContext
+    ) -> List[str]:
+        context.register_import("numpy", "np")
+        func_name = type(op).__name__.lower()[6:]
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        kw_str = self.generate_call_args_with_attributes(
+            op, context, "n", kw_keys=["d"]
+        )
+        return [f"{res_var_name} = np.fft.{func_name}({kw_str})"]

maxframe/codegen/spe/tensor/linalg.py

@@ -15,9 +15,12 @@
 from typing import List
 
 from ....tensor.linalg import (
+    TensorCholesky,
+    TensorEinsum,
     TensorInv,
     TensorLU,
     TensorMatmul,
+    TensorNorm,
     TensorQR,
     TensorSolveTriangular,
     TensorSVD,
@@ -26,17 +29,22 @@ from ....tensor.linalg import (
 from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
 from ..utils import build_method_call_adapter
 
+TensorCholeskyAdapter = build_method_call_adapter(
+    TensorCholesky, "cholesky", 0, kw_keys=["lower"], source_module="npl"
+)
 TensorInvAdapter = build_method_call_adapter(
     TensorInv, "inv", "input", source_module="npl"
 )
 TensorLUAdapter = build_method_call_adapter(
     TensorLU, "lu", "input", source_module="spl"
 )
-TensorQRAdapter = build_method_call_adapter(TensorQR, "qr", 0, source_module="spl")
-TensorSVDAdapter = build_method_call_adapter(TensorSVD, "svd", 0, source_module="spl")
 TensorMatmulAdapter = build_method_call_adapter(
     TensorMatmul, "matmul", "a", "b", source_module="np", kw_keys=["casting", "order"]
 )
+TensorNormAdapter = build_method_call_adapter(
+    TensorNorm, "norm", 0, kw_keys=["ord", "axis", "keepdims"], source_module="npl"
+)
+TensorQRAdapter = build_method_call_adapter(TensorQR, "qr", 0, source_module="spl")
 TensorSolveTriangularAdapter = build_method_call_adapter(
     TensorSolveTriangular,
     "solve_triangular",
@@ -45,6 +53,25 @@ TensorSolveTriangularAdapter = build_method_call_adapter(
     source_module="spl",
     kw_keys=["lower", "strict"],
 )
+TensorSVDAdapter = build_method_call_adapter(TensorSVD, "svd", 0, source_module="spl")
+
+
+@register_op_adapter(TensorEinsum)
+class TensorEinsumAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: TensorEinsum, context: SPECodeContext) -> List[str]:
+        context.register_import("numpy", "np")
+
+        inp_var_names_str = ", ".join(
+            [context.get_input_tileable_variable(inp) for inp in op.inputs]
+        )
+        res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        subscripts_str = self.translate_var(context, op.subscripts)
+        kw_str = self.generate_call_args_with_attributes(
+            op, context, kw_keys=["dtype", "order", "casting"]
+        )
+        return [
+            f"{res_var_name} = np.einsum({subscripts_str}, {inp_var_names_str}, {kw_str})"
+        ]
 
 
 @register_op_adapter(TensorTensorDot)

maxframe/codegen/spe/tensor/misc.py

@@ -16,11 +16,17 @@ from typing import List
 
 import numpy as np
 
+from ....tensor.indexing import TensorFillDiagonal
 from ....tensor.misc import (
+    TensorArgwhere,
+    TensorCopyTo,
+    TensorDelete,
     TensorDiff,
+    TensorGetShape,
     TensorIsIn,
     TensorRepeat,
     TensorSearchsorted,
+    TensorSplit,
     TensorSqueeze,
     TensorSwapAxes,
     TensorTranspose,
@@ -28,15 +34,39 @@ from ....tensor.misc import (
     TensorUnique,
     TensorWhere,
 )
+from ....tensor.misc.insert import TensorInsert
 from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
 from ..utils import build_method_call_adapter
 
 _trapz_func = "trapezoid" if hasattr(np, "trapezoid") else "trapz"
 
 
+TensorArgwhereAdapter = build_method_call_adapter(
+    TensorArgwhere, "argwhere", 0, source_module="np"
+)
+TensorDeleteAdapter = build_method_call_adapter(
+    TensorDelete, "delete", 0, "index_obj", kw_keys=["axis"], source_module="np"
+)
 TensorDiffAdapter = build_method_call_adapter(
     TensorDiff, "diff", 0, skip_none=True, kw_keys=["n", "axis"], source_module="np"
 )
+TensorFillDiagonalAdapter = build_method_call_adapter(
+    TensorFillDiagonal,
+    "fill_diagonal",
+    0,
+    "val",
+    skip_none=True,
+    kw_keys=["wrap"],
+    source_module="np",
+)
+TensorInsertAdapter = build_method_call_adapter(
+    TensorInsert,
+    "insert",
+    0,
+    "index_obj",
+    kw_keys=["values", "axis"],
+    source_module="np",
+)
 TensorIsInAdapter = build_method_call_adapter(
     TensorIsIn,
     "isin",
@@ -62,18 +92,6 @@ TensorSearchsortedAdapter = build_method_call_adapter(
     kw_keys=["side", "sorter"],
     source_module="np",
 )
-
-
-@register_op_adapter(TensorTranspose)
-class TensorTransposeAdapter(SPEOperatorAdapter):
-    def generate_code(self, op: TensorTranspose, context: SPECodeContext) -> List[str]:
-        context.register_import("numpy", "np")
-        input_var_name = context.get_input_tileable_variable(op.inputs[0])
-        res_var_name = context.get_output_tileable_variable(op.outputs[0])
-        axes_var = self.translate_var(context, op.axes)
-        return [f"{res_var_name} = np.transpose({input_var_name}, axes={axes_var})"]
-
-
 TensorTrapezoidAdapter = build_method_call_adapter(
     TensorTrapezoid,
     _trapz_func,
@@ -81,16 +99,60 @@ TensorTrapezoidAdapter = build_method_call_adapter(
     kw_keys=["y", "x", "dx", "axis"],
     source_module="np",
 )
+TensorSqueezeAdapter = build_method_call_adapter(
+    TensorSqueeze, "squeeze", 0, kw_keys=["axis"], source_module="np"
+)
+TensorSwapAxesAdapter = build_method_call_adapter(
+    TensorSwapAxes, "swapaxes", "axis1", "axis2", source_module="np"
+)
+TensorTransposeAdapter = build_method_call_adapter(
+    TensorTranspose, "transpose", 0, kw_keys=["axes"], source_module="np"
+)
+TensorWhereAdapter = build_method_call_adapter(
+    TensorWhere, "where", "condition", "x", "y", source_module="np"
+)
 
 
-@register_op_adapter(TensorSqueeze)
-class TensorSqueezeAdapter(SPEOperatorAdapter):
-    def generate_code(self, op: TensorSqueeze, context: SPECodeContext) -> List[str]:
+@register_op_adapter(TensorCopyTo)
+class TensorCopyToAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: TensorCopyTo, context: SPECodeContext) -> List[str]:
         context.register_import("numpy", "np")
-        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        dst_var_name = context.get_input_tileable_variable(op.inputs[1])
+        src_var_name = context.get_input_tileable_variable(op.inputs[0])
         res_var_name = context.get_output_tileable_variable(op.outputs[0])
+        args_str = self.generate_call_args_with_attributes(
+            op, context, kw_keys=["casting", "where"], skip_none=True
+        )
+        return [
+            f"{res_var_name} = {dst_var_name}.copy()",
+            f"np.copyto({res_var_name}, {src_var_name}, {args_str})",
+        ]
+
+
+@register_op_adapter(TensorGetShape)
+class TensorGetShapeAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: TensorGetShape, context: SPECodeContext) -> List[str]:
+        context.register_import("numpy", "np")
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_names = ", ".join(
+            context.get_output_tileable_variable(out) for out in op.outputs
+        )
+        return [f"{res_var_names} = np.shape({input_var_name})"]
+
+
+@register_op_adapter(TensorSplit)
+class TensorSplitAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: TensorSplit, context: SPECodeContext) -> List[str]:
+        context.register_import("numpy", "np")
+        input_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_names = ", ".join(
+            context.get_output_tileable_variable(out) for out in op.outputs
+        )
+        idx_var = self.translate_var(context, op.indices_or_sections)
         axis_var = self.translate_var(context, op.axis)
-        return [f"{res_var_name} = np.squeeze({input_var_name}, axis={axis_var})"]
+        return [
+            f"{res_var_names} = np.split({input_var_name}, {idx_var} axis={axis_var})"
+        ]
 
 
 @register_op_adapter(TensorUnique)
@@ -111,11 +173,3 @@ class TensorUniqueAdapter(SPEOperatorAdapter):
             f"return_inverse={return_inverse_var}, return_counts={return_counts_var}, "
             f"axis={axis_var})"
         ]
-
-
-TensorSwapAxesAdapter = build_method_call_adapter(
-    TensorSwapAxes, "swapaxes", "axis1", "axis2", source_module="np"
-)
-TensorWhereAdapter = build_method_call_adapter(
-    TensorWhere, "where", "condition", "x", "y", source_module="np"
-)

maxframe/codegen/spe/tensor/spatial.py

@@ -0,0 +1,45 @@
+# 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 ....tensor.spatial.distance.cdist import TensorCDist
+from ....tensor.spatial.distance.pdist import TensorPDist
+from ....tensor.spatial.distance.squareform import TensorSquareForm
+from ..utils import build_method_call_adapter
+
+TensorCDistAdapter = build_method_call_adapter(
+    TensorCDist,
+    "cdist",
+    "xa",
+    "xb",
+    kw_keys=["metric", "p", "w"],
+    v="V",
+    vi="VI",
+    source_module="scipy.spatial.distance",
+)
+TensorPDistAdapter = build_method_call_adapter(
+    TensorPDist,
+    "pdist",
+    0,
+    kw_keys=["metric", "p", "w"],
+    v="V",
+    vi="VI",
+    source_module="scipy.spatial.distance",
+)
+TensorSquareFormAdapter = build_method_call_adapter(
+    TensorSquareForm,
+    "squareform",
+    0,
+    kw_keys=["checks"],
+    source_module="scipy.spatial.distance",
+)

maxframe/codegen/spe/tensor/statistics.py

@@ -12,7 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List
+
 from ....tensor.statistics.bincount import TensorBinCount
+from ....tensor.statistics.digitize import TensorDigitize
+from ....tensor.statistics.histogram import TensorHistogram, TensorHistogramBinEdges
+from ....tensor.statistics.quantile import TensorQuantile
+from ..core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
 from ..utils import build_method_call_adapter
 
 TensorBinCountAdapter = build_method_call_adapter(
@@ -22,3 +28,41 @@ TensorBinCountAdapter = build_method_call_adapter(
     kw_keys=["weights", "minlength"],
     source_module="np",
 )
+TensorDigitizeAdapter = build_method_call_adapter(
+    TensorDigitize,
+    "digitize",
+    0,
+    kw_keys=["bins", "right"],
+    source_module="np",
+)
+TensorHistogramBinEdgesAdapter = build_method_call_adapter(
+    TensorHistogramBinEdges,
+    "histogram_bin_edges",
+    0,
+    kw_keys=["bins", "range", "weights"],
+    source_module="np",
+)
+TensorQuantileAdapter = build_method_call_adapter(
+    TensorQuantile,
+    "quantile",
+    0,
+    kw_keys=["q", "axis", "overwrite_input", "interpolation", "keepdims"],
+    source_module="np",
+)
+
+
+@register_op_adapter(TensorHistogram)
+class TensorHistogramAdapter(SPEOperatorAdapter):
+    def generate_code(self, op: TensorHistogram, context: SPECodeContext) -> List[str]:
+        context.register_import("numpy", "np")
+
+        inp_var_name = context.get_input_tileable_variable(op.inputs[0])
+        res_var_names = ", ".join(
+            context.get_output_tileable_variable(out) for out in op.outputs
+        )
+        if not op.ret_bins:
+            res_var_names += ", _"
+        kw_str = self.generate_call_args_with_attributes(
+            op, context, kw_keys=["bins", "range", "density", "weights"]
+        )
+        return [f"{res_var_names} = np.histogram({inp_var_name}, {kw_str})"]

maxframe/codegen/spe/tensor/tests/test_fft.py

@@ -0,0 +1,64 @@
+# 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 maxframe.tensor as mt
+
+from ...core import SPECodeContext
+from ..fft import (
+    TensorFFTAdapter,
+    TensorFFTFreqAdapter,
+    TensorFFTNAdapter,
+    TensorFFTShiftAdapter,
+)
+
+
+def test_fft():
+    result = mt.fft.fft(mt.exp(2j * mt.pi * mt.arange(8) / 8))
+
+    adapter = TensorFFTAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = ["var_1 = np.fft.fft(var_0, axis=-1)"]
+    assert results == expected_results
+
+
+def test_fft2():
+    a = mt.mgrid[:5, :5][0]
+    result = mt.fft.fft2(a)
+
+    adapter = TensorFFTNAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = ["var_1 = np.fft.fft2(var_0, None, axes=(-2, -1))"]
+    assert results == expected_results
+
+
+def test_fftshift():
+    result = mt.fft.fftshift(mt.tensor(10), 0.1)
+
+    adapter = TensorFFTShiftAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = ["var_1 = np.fft.fftshift(var_0, axes=(0.1,))"]
+    assert results == expected_results
+
+
+def test_fftfreq():
+    result = mt.fft.fftfreq(10, 0.1)
+
+    adapter = TensorFFTFreqAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = ["var_0 = np.fft.fftfreq(10, d=0.1)"]
+    assert results == expected_results

maxframe/codegen/spe/tensor/tests/test_linalg.py

@@ -14,7 +14,21 @@
 
 from ..... import tensor as mt
 from ...core import SPECodeContext
-from ..linalg import TensorQRAdapter, TensorTensorDotAdapter
+from ..linalg import TensorEinsumAdapter, TensorQRAdapter, TensorTensorDotAdapter
+
+
+def test_einsum():
+    a = mt.arange(25).reshape(5, 5)
+    b = mt.arange(5)
+    adapter = TensorEinsumAdapter()
+    r = mt.einsum("ij,j", a, b)
+    context = SPECodeContext()
+    results = adapter.generate_code(r.op, context)
+    expected_results = [
+        "var_2 = np.einsum('ij,j->i', var_0, var_1, dtype=np.dtype('int64'),"
+        " order='K', casting='safe')"
+    ]
+    assert results == expected_results
 
 
 def test_qr():

maxframe/codegen/spe/tensor/tests/test_misc.py

@@ -15,7 +15,11 @@
 from ..... import tensor as mt
 from ...core import SPECodeContext
 from ..misc import (
+    TensorCopyToAdapter,
+    TensorGetShapeAdapter,
+    TensorInsertAdapter,
     TensorIsInAdapter,
+    TensorSplitAdapter,
     TensorSqueezeAdapter,
     TensorTransposeAdapter,
     TensorUniqueAdapter,
@@ -40,7 +44,7 @@ def test_transpose():
     adapter = TensorTransposeAdapter()
     context = SPECodeContext()
     results = adapter.generate_code(result.op, context)
-    expected_results = ["var_1 = np.transpose(var_0, axes=[1, 0, 2])"]
+    expected_results = ["var_0 = np.transpose(var_1, axes=[1, 0, 2])"]
     assert results == expected_results
 
 
@@ -50,7 +54,7 @@ def test_squeeze():
     adapter = TensorSqueezeAdapter()
     context = SPECodeContext()
     results = adapter.generate_code(result.op, context)
-    expected_results = ["var_1 = np.squeeze(var_0, axis=(0,))"]
+    expected_results = ["var_0 = np.squeeze(var_1, axis=(0,))"]
     assert results == expected_results
 
 
@@ -92,3 +96,49 @@ def test_unique():
         "np.unique(var_0, return_index=True, return_inverse=True, return_counts=True, axis=0)"
     ]
     assert results == expected_results
+
+
+def test_insert():
+    result = mt.insert(mt.array([[1, 1], [2, 2], [3, 3]]), 1, 5)
+
+    adapter = TensorInsertAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = ["var_0 = np.insert(var_1, 1, values=5)"]
+    assert results == expected_results
+
+
+def test_split():
+    res_tp = mt.split(mt.array([1, 2, 3, 4, 5, 6]), 3)
+
+    adapter = TensorSplitAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(res_tp[0].op, context)
+    expected_results = ["var_1, var_2, var_3 = np.split(var_0, 3 axis=0)"]
+    assert results == expected_results
+
+
+def test_get_shape():
+    inp = mt.array([[1, 1], [2, 2], [3, 3]])
+    res = mt.shape(inp[:, inp[0, :] == 1])
+
+    adapter = TensorGetShapeAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(res[0].op, context)
+    expected_results = ["var_1, var_2 = np.shape(var_0)"]
+    assert results == expected_results
+
+
+def test_copyto():
+    dest = mt.array([[1, 1], [2, 2], [3, 3]])
+    src = mt.array([[1, 1], [2, 2], [3, 4]])
+    mt.copyto(dest, src)
+
+    adapter = TensorCopyToAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(dest.op, context)
+    expected_results = [
+        "var_2 = var_0.copy()",
+        "np.copyto(var_2, var_1, casting='same_kind')",
+    ]
+    assert results == expected_results

maxframe/codegen/spe/tensor/tests/test_spatial.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 ..... import tensor as mt
+from ...core import SPECodeContext
+from ..spatial import TensorCDistAdapter
+
+
+def test_cdist():
+    xa = mt.random.rand(100, 10)
+    xb = mt.random.rand(89, 10)
+    VI = mt.random.rand(10, 10)
+    dist = mt.spatial.distance.cdist(xa, xb, metric="mahalanobis", VI=VI)
+
+    adapter = TensorCDistAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(dist.op, context)
+    expected_results = [
+        "var_0 = scipy.spatial.distance.cdist("
+        "var_1, var_2, metric='mahalanobis', VI=var_3)"
+    ]
+    assert results == expected_results

maxframe/codegen/spe/tensor/tests/test_statistics.py

@@ -14,7 +14,7 @@
 
 from ..... import tensor as mt
 from ...core import SPECodeContext
-from ..statistics import TensorBinCountAdapter
+from ..statistics import TensorBinCountAdapter, TensorHistogramAdapter
 
 
 def test_bincount():
@@ -27,3 +27,17 @@ def test_bincount():
     results = adapter.generate_code(result.op, context)
     expected_results = ["var_0 = np.bincount(var_1, weights=var_2, minlength=0)"]
     assert results == expected_results
+
+
+def test_histogram():
+    arr = mt.arange(5)
+    result, _bins = mt.histogram(arr, bins=3)
+
+    adapter = TensorHistogramAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = [
+        "var_1, _ = np.histogram(var_0, bins=var_2, range=None, density=None, "
+        "weights=None)"
+    ]
+    assert results == expected_results

maxframe/codegen/spe/tests/test_spe_codegen.py

@@ -59,12 +59,8 @@ def test_codegen_with_udf(codegen):
 
     udf_1_body = base64.b64encode(pickle.dumps(f1, protocol=pickle.DEFAULT_PROTOCOL))
     udf_2_body = base64.b64encode(pickle.dumps(f2, protocol=pickle.DEFAULT_PROTOCOL))
-    udf_1_value = (
-        f"pickled_data = cloudpickle.loads(base64.b64decode({udf_1_body}), buffers=[])"
-    )
-    udf_2_value = (
-        f"pickled_data = cloudpickle.loads(base64.b64decode({udf_2_body}), buffers=[])"
-    )
+    udf_1_value = f"udf_main_entry = cloudpickle.loads(base64.b64decode({udf_1_body}), buffers=[])"
+    udf_2_value = f"udf_main_entry = cloudpickle.loads(base64.b64decode({udf_2_body}), buffers=[])"
     udf_1 = f"user_udf_f1_{hashlib.md5(udf_1_value.encode('utf-8')).hexdigest()}"
     udf_2 = f"user_udf_f2_{hashlib.md5(udf_2_value.encode('utf-8')).hexdigest()}"
 
@@ -74,9 +70,9 @@ def test_codegen_with_udf(codegen):
     import base64
     import cloudpickle
     {udf_1_value}
-    {udf_1} = pickled_data
+    {udf_1} = udf_main_entry
     {udf_2_value}
-    {udf_2} = pickled_data
+    {udf_2} = udf_main_entry
     if not running:
         raise RuntimeError('CANCELLED')
     var_0 = np.random.rand(1, 3)
@@ -111,9 +107,7 @@ def test_codegen_with_udf_and_args(codegen):
             protocol=pickle.DEFAULT_PROTOCOL,
         )
     )
-    udf_1_value = (
-        f"pickled_data = cloudpickle.loads(base64.b64decode({udf_1_body}), buffers=[])"
-    )
+    udf_1_value = f"udf_main_entry = cloudpickle.loads(base64.b64decode({udf_1_body}), buffers=[])"
     udf_1 = f"user_udf_f1_{hashlib.md5(udf_1_value.encode('utf-8')).hexdigest()}"
 
     expected_contents = f"""
@@ -122,7 +116,7 @@ def test_codegen_with_udf_and_args(codegen):
     import base64
     import cloudpickle
     {udf_1_value}
-    {udf_1} = pickled_data
+    {udf_1} = udf_main_entry
     if not running:
         raise RuntimeError('CANCELLED')
     var_0 = np.random.rand(1, 3)

maxframe/codegen/spe/utils.py

@@ -37,6 +37,8 @@ def build_method_call_adapter(
         def generate_code(self, op: OperatorType, context: SPECodeContext) -> List[str]:
             if source_module in _import_aliases:
                 context.register_import(_import_aliases[source_module], source_module)
+            elif source_module:
+                context.register_import(source_module)
 
             input_var_name = source_module or context.get_input_tileable_variable(
                 op.inputs[0]