maxframe 2.4.0rc1__cp312-cp312-win32.whl

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

@@ -0,0 +1,39 @@
+# 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 ..reshape import TensorReshapeAdapter
+
+
+def test_reshape_C():
+    t1 = mt.arange(6).reshape((3, 2))
+    t2 = mt.reshape(t1, (2, 3))
+    adapter = TensorReshapeAdapter()
+    context = SPECodeContext()
+    context.get_input_tileable_variable(t1.op.outputs[0])
+    results = adapter.generate_code(t2.op, context)
+    expected_results = ["var_1 = np.reshape(var_0, (2, 3), 'C')"]
+    assert results == expected_results
+
+
+def test_reshape_F():
+    t1 = mt.arange(6).reshape((3, 2))
+    t2 = mt.reshape(t1, (2, 3), "F")
+    adapter = TensorReshapeAdapter()
+    context = SPECodeContext()
+    context.get_input_tileable_variable(t1.op.outputs[0])
+    results = adapter.generate_code(t2.op, context)
+    expected_results = ["var_1 = np.reshape(var_0, (2, 3), 'F')"]
+    assert results == expected_results

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

@@ -0,0 +1,49 @@
+# 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 ..sort import TensorSortAdapter
+
+
+def test_sort():
+    t1 = mt.array([1, 4, 5, 3])
+    t2 = mt.sort(t1, stable=True)
+    adapter = TensorSortAdapter()
+    context = SPECodeContext()
+    context.get_input_tileable_variable(t1.op.outputs[0])
+    results = adapter.generate_code(t2.op, context)
+    expected_results = ["var_1 = np.sort(var_0, axis=0, kind='STABLE', order=None)"]
+    assert results == expected_results
+
+    t2, t3 = mt.sort(t1, stable=True, return_index=True)
+    adapter = TensorSortAdapter()
+    context = SPECodeContext()
+    context.get_input_tileable_variable(t1.op.outputs[0])
+    results = adapter.generate_code(t2.op, context)
+    expected_results = [
+        "var_1 = np.sort(var_0, axis=0, kind='STABLE', order=None)",
+        "var_2 = np.argsort(var_0, axis=0, kind='STABLE', order=None)",
+    ]
+    assert results == expected_results
+
+    t2 = mt.argsort(t1, stable=True)
+    adapter = TensorSortAdapter()
+    context = SPECodeContext()
+    context.get_input_tileable_variable(t1.op.outputs[0])
+    results = adapter.generate_code(t2.op, context)
+    expected_results = [
+        "var_1 = np.argsort(var_0, axis=0, kind='STABLE', order=None)",
+    ]
+    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_special.py

@@ -0,0 +1,28 @@
+# 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 .....tensor.special import expit
+from ...core import SPECodeContext
+from ..special import TensorSpecialUnaryOpAdapter
+
+
+def test_special_unary():
+    result = expit(mt.random.rand(10, 10))
+
+    adapter = TensorSpecialUnaryOpAdapter()
+    context = SPECodeContext()
+    results = adapter.generate_code(result.op, context)
+    expected_results = ["var_1 = spspecial.expit(var_0)"]
+    assert results == expected_results

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

@@ -0,0 +1,43 @@
+# 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 ..statistics import TensorBinCountAdapter, TensorHistogramAdapter
+
+
+def test_bincount():
+    w = mt.array([0.3, 0.5, 0.2, 0.7, 1.0, -0.6])  # weights
+    x = mt.array([0, 1, 1, 2, 2, 2])
+    result = mt.bincount(x, weights=w)
+
+    adapter = TensorBinCountAdapter()
+    context = SPECodeContext()
+    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/__init__.py

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

maxframe/codegen/spe/tests/test_remote.py

@@ -0,0 +1,29 @@
+# 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 remote as mr
+from ..core import SPECodeContext
+from ..remote import RemoteFunctionAdapter
+
+
+def test_remote():
+    def f(v1, v2):
+        return v1 + v2
+
+    v0 = mr.spawn(f, (2, 3))
+    v1 = mr.spawn(f, (v0, 10))
+    context = SPECodeContext()
+    results = RemoteFunctionAdapter().generate_code(v1.op, context)
+    assert callable(context.constants["const_0"])
+    assert results[0] == "var_1 = const_0(var_0, 10)"

maxframe/codegen/spe/tests/test_spe_codegen.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.
+
+import base64
+import hashlib
+import textwrap
+
+import pytest
+
+from .... import dataframe as md
+from .... import tensor as mt
+from ....dataframe.utils import bind_func_args_from_pos
+from ....lib import wrapped_pickle as pickle
+from ..core import SPECodeGenerator
+
+
+@pytest.fixture
+def codegen():
+    return SPECodeGenerator("session_1")
+
+
+def test_simple_codegen(codegen):
+    df = md.DataFrame(mt.random.rand(1000, 5), columns=list("ABCDE"))
+    df["F"] = df["A"] + df["B"]
+    df["G"] = df["C"] * 2 + df["D"]
+    result = df.cumsum()
+    dag = result.build_graph()
+    generated = codegen.generate(dag)
+
+    local_vars = generated.constants.copy()
+    exec(generated.code, {}, local_vars)
+    df_result = local_vars[generated.output_key_to_variables[result.key]]
+    assert len(df_result.columns) == 7
+
+
+def test_codegen_with_udf(codegen):
+    df = md.DataFrame(mt.random.rand(1, 3), columns=list("ABC"))
+
+    def f1(x) -> int:
+        return x + 1
+
+    def f2(x) -> int:
+        return x - 1
+
+    result = df.transform({"A": f1, "B": [f1, f2]})
+    dag = result.build_graph()
+    generated = codegen.generate(dag)
+
+    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"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()}"
+
+    expected_contents = f"""
+    import numpy as np
+    import pandas as pd
+    import base64
+    import cloudpickle
+    {udf_1_value}
+    {udf_1} = udf_main_entry
+    {udf_2_value}
+    {udf_2} = udf_main_entry
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_0 = np.random.rand(1, 3)
+
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_1 = pd.DataFrame(var_0, index=const_0, columns=['A', 'B', 'C'])
+
+    del var_0
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_2 = var_1.transform({{'A': {udf_1}, 'B': [{udf_1}, {udf_2}]}}, axis=0)
+
+    del var_1
+    """
+    assert generated.code == textwrap.dedent(expected_contents).strip()
+
+
+def test_codegen_with_udf_and_args(codegen):
+    df = md.DataFrame(mt.random.rand(1, 3), columns=list("ABC"))
+
+    def f1(x, y, a, b) -> int:
+        return x + y + a + b
+
+    result = df.transform({"A": f1}, 0, 2, a=3, b=4)
+    dag = result.build_graph()
+    generated = codegen.generate(dag)
+
+    udf_1_body = base64.b64encode(
+        pickle.dumps(
+            bind_func_args_from_pos(f1, 1, 2, a=3, b=4),
+            protocol=pickle.DEFAULT_PROTOCOL,
+        )
+    )
+    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"""
+    import numpy as np
+    import pandas as pd
+    import base64
+    import cloudpickle
+    {udf_1_value}
+    {udf_1} = udf_main_entry
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_0 = np.random.rand(1, 3)
+
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_1 = pd.DataFrame(var_0, index=const_0, columns=['A', 'B', 'C'])
+
+    del var_0
+    if not running:
+        raise RuntimeError('CANCELLED')
+    var_2 = var_1.transform({{'A': {udf_1}}}, axis=0)
+
+    del var_1
+    """
+    assert generated.code == textwrap.dedent(expected_contents).strip()

maxframe/codegen/spe/utils.py

@@ -0,0 +1,56 @@
+# 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, Type
+
+from ...typing_ import OperatorType
+from .core import SPECodeContext, SPEOperatorAdapter, register_op_adapter
+
+_import_aliases = {
+    "np": "numpy",
+    "npl": "numpy.linalg",
+    "pd": "pandas",
+    "spl": "scipy.linalg",
+}
+
+
+def build_method_call_adapter(
+    op_cls: Type[OperatorType],
+    method_name: str,
+    *call_args,
+    skip_none: bool = True,
+    source_module: str = None,
+    **call_kwargs,
+) -> Type[SPEOperatorAdapter]:
+    class MethodCallAdapter(SPEOperatorAdapter):
+        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]
+            )
+            res_var_names = [
+                context.get_output_tileable_variable(out) for out in op.outputs
+            ]
+            args_str = self.generate_call_args_with_attributes(
+                op, context, *call_args, skip_none=skip_none, **call_kwargs
+            )
+            res_var_names_str = ", ".join(res_var_names)
+            return [f"{res_var_names_str} = {input_var_name}.{method_name}({args_str})"]
+
+    MethodCallAdapter.__name__ = op_cls.__name__ + "Adapter"
+    return register_op_adapter(op_cls)(MethodCallAdapter)

maxframe/codegen/tests/__init__.py

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

maxframe/codegen/tests/test_codegen.py

@@ -0,0 +1,67 @@
+# 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 base64
+from typing import List, Tuple
+
+import pytest
+
+from ...lib import wrapped_pickle
+from ...serialization.core import PickleContainer
+from ..core import UserCodeMixin
+
+
+@pytest.mark.parametrize(
+    "input_obj, expected_output",
+    [
+        (None, "None"),
+        (10, "10"),
+        (3.14, "3.14"),
+        (True, "True"),
+        (False, "False"),
+        (b"hello", "base64.b64decode(b'aGVsbG8=')"),
+        ("hello", "'hello'"),
+        ([1, 2, 3], "[1, 2, 3]"),
+        ({"a": 1, "b": 2}, "{'a': 1, 'b': 2}"),
+        ((1, 2, 3), "(1, 2, 3)"),
+        ((1,), "(1,)"),
+        ((), "()"),
+        ({1, 2, 3}, "{1, 2, 3}"),
+        (set(), "set()"),
+    ],
+)
+def test_obj_to_python_expr(input_obj, expected_output):
+    assert UserCodeMixin.obj_to_python_expr(input_obj) == expected_output
+
+
+def test_obj_to_python_expr_custom_object():
+    class CustomClass:
+        def __init__(self, a: int, b: List[int], c: Tuple[int, int]):
+            self.a = a
+            self.b = b
+            self.c = c
+
+    custom_obj = CustomClass(1, [2, 3], (4, 5))
+    pickle_data = wrapped_pickle.dumps(custom_obj)
+    pickle_str = base64.b64encode(pickle_data)
+    custom_obj_pickle_container = PickleContainer([pickle_data])
+
+    # with class obj will not support currently
+    with pytest.raises(ValueError):
+        UserCodeMixin.obj_to_python_expr(custom_obj)
+
+    assert (
+        UserCodeMixin.obj_to_python_expr(custom_obj_pickle_container)
+        == f"cloudpickle.loads(base64.b64decode({pickle_str}), buffers=[])"
+    )

maxframe/config/__init__.py

@@ -0,0 +1,15 @@
+# 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 .config import AttributeDict, option_context, options