maxframe 2.0.0b2__cp37-cp37m-win32.whl → 2.3.0rc1__cp37-cp37m-win32.whl

maxframe/codegen/spe/learn/metrics/pairwise.py

@@ -0,0 +1,51 @@
+# 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 .....learn.metrics.pairwise.cosine import CosineDistances
+from .....learn.metrics.pairwise.euclidean import EuclideanDistances
+from .....learn.metrics.pairwise.haversine import HaversineDistances
+from .....learn.metrics.pairwise.manhattan import ManhattanDistances
+from ...utils import build_method_call_adapter
+
+CosineDistancesAdapter = build_method_call_adapter(
+    CosineDistances,
+    "cosine_distances",
+    "x",
+    "y",
+    source_module="sklearn.metrics.pairwise",
+)
+EuclideanDistancesAdapter = build_method_call_adapter(
+    EuclideanDistances,
+    "euclidean_distances",
+    "x",
+    "y",
+    y_norm_squared="Y_norm_squared",
+    squared="squared",
+    x_norm_squared="X_norm_squared",
+    source_module="sklearn.metrics.pairwise",
+)
+HaversineDistancesAdapter = build_method_call_adapter(
+    HaversineDistances,
+    "haversine_distances",
+    "x",
+    "y",
+    source_module="sklearn.metrics.pairwise",
+)
+ManhattanDistancesAdapter = build_method_call_adapter(
+    ManhattanDistances,
+    "manhattan_distances",
+    "x",
+    "y",
+    source_module="sklearn.metrics.pairwise",
+)

maxframe/codegen/spe/learn/metrics/tests/test_pairwise.py

@@ -0,0 +1,36 @@
+# 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 ......learn.metrics.pairwise import euclidean_distances
+from ....core import SPECodeContext
+from ..pairwise import EuclideanDistancesAdapter
+
+
+def test_euclidean_distances():
+    x = mt.random.rand(10, 3)
+
+    d = euclidean_distances(
+        x,
+        X_norm_squared=mt.random.rand(10, 1),
+    )
+
+    adapter = EuclideanDistancesAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(d.op, context)
+    expected_results = [
+        "var_0 = sklearn.metrics.pairwise.euclidean_distances("
+        "var_1, var_1, squared=False, X_norm_squared=var_2)"
+    ]
+    assert results == expected_results

maxframe/codegen/spe/learn/metrics/tests/test_ranking.py

@@ -0,0 +1,59 @@
+# 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 ......learn.metrics import auc, roc_auc_score, roc_curve
+from ....core import SPECodeContext
+from .._ranking import AucOpAdapter, RocAucScoreAdapter, RocCurveAdapter
+
+
+def test_auc():
+    t1 = mt.array([1, 1, 2, 2])
+    t2 = mt.array([0.1, 0.4, 0.35, 0.8])
+    result = auc(t1, t2)
+
+    adapter = AucOpAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = ["var_2 = sk_auc(var_0, var_1)"]
+    assert results == expected_results
+
+
+def test_roc_auc_score():
+    t1 = mt.array([0, 0, 1, 1])
+    t2 = mt.array([0.1, 0.4, 0.35, 0.8])
+    result = roc_auc_score(t1, t2)
+
+    adapter = RocAucScoreAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = [
+        "var_2 = sk_roc_auc_score("
+        "var_0, var_1, average='macro', multi_class='raise')"
+    ]
+    assert results == expected_results
+
+
+def test_roc_curve():
+    t1 = mt.array([0, 0, 1, 1])
+    t2 = mt.array([0.1, 0.4, 0.35, 0.8])
+    result = roc_curve(t1, t2)
+
+    adapter = RocCurveAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result[0].op, context)
+    expected_results = [
+        "var_2, var_3, var_4 = sk_roc_curve(" "var_0, var_1, drop_intermediate=True)"
+    ]
+    assert results == expected_results

maxframe/codegen/spe/tensor/__init__.py

@@ -17,6 +17,7 @@ from . import (
     datasource,
     extensions,
     fetch,
+    fft,
     indexing,
     linalg,
     merge,
@@ -24,5 +25,7 @@ from . import (
     random,
     reduction,
     reshape,
+    spatial,
     special,
+    statistics,
 )

maxframe/codegen/spe/tensor/datasource.py

@@ -142,6 +142,7 @@ class TensorOnesZerosLikeAdapter(SPEOperatorAdapter):
         ]
 
 
+@register_op_adapter(Scalar)
 class TensorScalarAdapter(SPEOperatorAdapter):
     def accepts(self, op: Scalar) -> EngineAcceptance:
         return EngineAcceptance.SUCCESSOR

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():