mplang-nightly 0.1.dev141__py3-none-any.whl → 0.1.dev143__py3-none-any.whl

mplang/frontend/phe.py

@@ -52,10 +52,9 @@ def add(operand1: TensorType, operand2: TensorType) -> TensorType:
 @_PHE_MOD.simple_op()
 def mul(ciphertext: TensorType, plaintext: TensorType) -> TensorType:
     """Multiply a PHE ciphertext with a plaintext value (ciphertext dtype preserved)."""
-    if ciphertext.dtype.is_floating and plaintext.dtype.is_floating:
+    if plaintext.dtype.is_floating:
         raise ValueError(
-            "PHE multiplication does not support float x float operations due to truncation requirements. "
-            "Consider using mixed types (float x int) or integer types instead."
+            "PHE multiplication does not support floating-point plaintext."
         )
     return ciphertext
 
@@ -65,3 +64,141 @@ def decrypt(ciphertext: TensorType, private_key: TensorType) -> TensorType:
     """Decrypt ciphertext using PHE private key: returns plaintext with same semantic type as ciphertext."""
     _ = private_key
     return ciphertext
+
+
+@_PHE_MOD.simple_op()
+def dot(ciphertext: TensorType, plaintext: TensorType) -> TensorType:
+    """Compute dot product of ciphertext with plaintext.
+
+    Args:
+        ciphertext: The ciphertext operand (first argument)
+        plaintext: The plaintext operand (second argument)
+
+    Returns:
+        TensorType: Result tensor type with computed shape following numpy dot product rules
+    """
+    # For dot product, we need to calculate the result shape
+    # This follows numpy dot product rules
+    import numpy as np
+
+    # Create dummy arrays to determine result shape
+    dummy_ct = np.zeros(ciphertext.shape)
+    dummy_pt = np.zeros(plaintext.shape)
+    dummy_result = np.dot(dummy_ct, dummy_pt)
+
+    return TensorType(ciphertext.dtype, dummy_result.shape)
+
+
+@_PHE_MOD.simple_op()
+def gather(ciphertext: TensorType, indices: TensorType, *, axis: int = 0) -> TensorType:
+    """Gather elements from ciphertext using indices.
+
+    Args:
+        ciphertext: The ciphertext to gather from
+        indices: The indices to gather
+        axis: The axis along which to gather (default: 0)
+    """
+    # Calculate result shape based on axis parameter
+    ct_shape = list(ciphertext.shape)
+    indices_shape = list(indices.shape)
+
+    # Normalize negative axis
+    normalized_axis = axis if axis >= 0 else len(ct_shape) + axis
+
+    # Result shape: replace the axis dimension with indices shape
+    result_shape = (
+        ct_shape[:normalized_axis] + indices_shape + ct_shape[normalized_axis + 1 :]
+    )
+    return TensorType(ciphertext.dtype, tuple(result_shape))
+
+
+@_PHE_MOD.simple_op()
+def scatter(
+    ciphertext: TensorType,
+    indices: TensorType,
+    updates: TensorType,
+    *,
+    axis: int = 0,
+) -> TensorType:
+    """Scatter updates into ciphertext at specified indices.
+
+    Args:
+        ciphertext: The ciphertext to scatter into
+        indices: The indices to scatter at
+        updates: The ciphertext updates to scatter
+        axis: The axis along which to scatter (default: 0)
+
+    Returns:
+        TensorType: Result tensor type with same shape and dtype as original ciphertext
+    """
+    return ciphertext
+
+
+@_PHE_MOD.simple_op()
+def concat(operand0: TensorType, operand1: TensorType, *, axis: int = 0) -> TensorType:
+    """Concatenate ciphertext tensors along specified axis.
+
+    Args:
+        operand0: The first ciphertext operand to concatenate
+        operand1: The second ciphertext operand to concatenate
+        axis: Axis along which to concatenate
+
+    Returns:
+        TensorType: Result tensor type with computed shape following numpy concatenation rules
+    """
+    # All operands should have same dtype
+    first_dtype = operand0.dtype
+    if operand1.dtype != first_dtype:
+        raise ValueError("All operands must have the same dtype for concatenation")
+
+    # Calculate result shape using numpy concatenation logic
+    import numpy as np
+
+    dummy_arrays = [np.zeros(operand0.shape), np.zeros(operand1.shape)]
+    dummy_result = np.concatenate(dummy_arrays, axis=axis)
+
+    return TensorType(first_dtype, dummy_result.shape)
+
+
+@_PHE_MOD.simple_op()
+def reshape(ciphertext: TensorType, *, new_shape: tuple[int, ...]) -> TensorType:
+    """Reshape ciphertext to new shape.
+
+    Args:
+        ciphertext: The ciphertext to reshape
+        new_shape: The target shape (can contain -1 for inferred dimension)
+
+    Returns:
+        TensorType: Result tensor type with computed shape following numpy reshape rules
+    """
+    # Calculate the actual result shape (handling -1 inference)
+    import numpy as np
+
+    dummy_array = np.zeros(ciphertext.shape)
+    # use this to check the correctness of new_shape
+    dummy_result = dummy_array.reshape(new_shape)
+    actual_shape = dummy_result.shape
+
+    return TensorType(ciphertext.dtype, actual_shape)
+
+
+@_PHE_MOD.simple_op()
+def transpose(
+    ciphertext: TensorType, *, axes: tuple[int, ...] | None = None
+) -> TensorType:
+    """Transpose ciphertext by permuting axes.
+
+    Args:
+        ciphertext: The ciphertext to transpose
+        axes: Permutation of axes (None for default reverse order)
+
+    Returns:
+        TensorType: Result tensor type with computed shape following numpy transpose rules
+    """
+    # Calculate result shape using numpy transpose logic
+    import numpy as np
+
+    dummy_array = np.zeros(ciphertext.shape)
+    dummy_result = np.transpose(dummy_array, axes)
+
+    return TensorType(ciphertext.dtype, dummy_result.shape)

{mplang_nightly-0.1.dev141.dist-info → mplang_nightly-0.1.dev143.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev141
+Version: 0.1.dev143
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev141.dist-info → mplang_nightly-0.1.dev143.dist-info}/RECORD

@@ -7,7 +7,7 @@ mplang/backend/__init__.py,sha256=Pn1MGW7FZ8ZQWcx_r2Io4Q1JbrMINCpefQt7yCNq_dc,74
 mplang/backend/base.py,sha256=zaofB1MPC9Af8FS-rx7q6utFxiu9Mppr2gWaKFZ70Os,9863
 mplang/backend/builtin.py,sha256=Mk1uUO2Vpw3meqZ0B7B0hG-wndea6cmFv2Uk1vM_uTg,7052
 mplang/backend/crypto.py,sha256=H_s5HI7lUP7g0xz-a9qMbSn6dhJStUilKbn3-7SIh0I,3812
-mplang/backend/phe.py,sha256=Y07fkTHTSKHgpMQ1soA3d2eBMlqr25uHCAtZOd5D5_I,10555
+mplang/backend/phe.py,sha256=uNqmrbDAbd97TWS_O6D5sopastHy6J20R7knFE4M4uc,65247
 mplang/backend/spu.py,sha256=i6Qqgeg-Anwpb5xX5Uz8GdmTWNkRy_pjp-xptIvlxl4,9273
 mplang/backend/sql_duckdb.py,sha256=WFEiJk3t68DMpjV9xZz5bJj3v8UMx4X5QWleLIE4Aq0,1499
 mplang/backend/stablehlo.py,sha256=GOxy-qgOxyEdtBkt6LASKzaZnPewZhvHYSPOgFFXgIM,2612
@@ -41,7 +41,7 @@ mplang/frontend/builtin.py,sha256=8qrlbe_SSy6QTXTnMG6_ADB8jSklVZGFBrkoR-p02FE,93
 mplang/frontend/crypto.py,sha256=Nf8zT4Eko7MIs4R2tgZecKVd7d6Hvd_CGGmANhs3Ghs,3651
 mplang/frontend/ibis_cc.py,sha256=01joUdFS_Ja9--PkezBhEcW_a9mkDrLgOhu5320s_bQ,4167
 mplang/frontend/jax_cc.py,sha256=ssP6rCvyWQ5VAr80-7z9QZUE2mWXyozJCGpq1dYQYY8,6374
-mplang/frontend/phe.py,sha256=zsL4rbukoZ4SvCkPdfgGQRzCwXk4_opn693RTgwp7gQ,2501
+mplang/frontend/phe.py,sha256=tDsCvStjVJ1Fs07yF3idkFnugUCA1zdFApPx7Uuulik,6795
 mplang/frontend/spu.py,sha256=yvnXH8HU55t7j_jaTpxa3Nbh0SqXufVDyZBTYCsRTK4,4994
 mplang/frontend/sql.py,sha256=DFdvjEPQX28VCRgUMeHYR0rwwOaoCH15bpvvlclLtHA,1999
 mplang/frontend/tee.py,sha256=EigmlbYDGvXkZCMHSYRAiOboSl9TG0ewoudbgl3_V6M,1393
@@ -69,8 +69,8 @@ mplang/utils/crypto.py,sha256=rvPomBFtznRHc3RPi6Aip9lsU8zW2oxBqGv1K3vn7Rs,1052
 mplang/utils/func_utils.py,sha256=vCJcZmu0bEbqhOQKdpttV2_MBllIcPSN0b8U4WjNGGo,5164
 mplang/utils/spu_utils.py,sha256=S3L9RBkBe2AvSuMSQQ12cBY5Y1NPthubvErSX_7nj1A,4158
 mplang/utils/table_utils.py,sha256=aC-IZOKkSmFkpr3NZchLM0Wt0GOn-rg_xHBHREWBwAU,2202
-mplang_nightly-0.1.dev141.dist-info/METADATA,sha256=k1IfwOhysZSd0aGI1mWuDSVFo8Hx3rwR69lkpZWb8Cc,16547
-mplang_nightly-0.1.dev141.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev141.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev141.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev141.dist-info/RECORD,,
+mplang_nightly-0.1.dev143.dist-info/METADATA,sha256=BvANiZEEU0N7uAY_k8Yd4p1rZKGPpdbrM5giBinXzdA,16547
+mplang_nightly-0.1.dev143.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev143.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev143.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev143.dist-info/RECORD,,