mplang-nightly 0.1.dev192__py3-none-any.whl → 0.1.dev268__py3-none-any.whl

mplang/v1/ops/crypto.py

@@ -0,0 +1,262 @@
+# Copyright 2025 Ant Group Co., 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.
+
+"""
+Crypto frontend operations: operation signatures, types, and high-level semantics.
+
+Scope and contracts:
+- This module defines portable API shapes; it does not implement cryptography.
+- Backends execute the operations and must meet the security semantics required
+    by the deployment (confidentiality, authenticity, correctness, etc.).
+- The enc/dec API in this frontend uses a conventional 12-byte nonce prefix
+    (ciphertext = nonce || payload), and dec expects that format. Other security
+    properties (e.g., AEAD) are backend responsibilities.
+"""
+
+from __future__ import annotations
+
+from jax.tree_util import PyTreeDef, tree_flatten
+
+from mplang.v1.core import UINT8, TensorType
+from mplang.v1.core.mpobject import MPObject
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.ops.base import stateless_mod
+
+_CRYPTO_MOD = stateless_mod("crypto")
+
+
+def _get_algo_overhead(algo: str) -> int:
+    """Get ciphertext overhead for a given encryption algorithm.
+
+    Args:
+        algo: Encryption algorithm identifier
+
+    Returns:
+        int: Number of overhead bytes added to plaintext length
+    """
+    overhead_map = {
+        "aes-ctr": 16,  # nonce only (legacy compatibility)
+        "aes-gcm": 28,  # nonce(12) + tag(16) for AES-GCM
+        "sm4-gcm": 28,  # nonce(12) + tag(16) for SM4-GCM
+    }
+
+    if algo not in overhead_map:
+        # return unknown overhead as -1
+        return -1
+    return overhead_map[algo]
+
+
+@_CRYPTO_MOD.simple_op()
+def keygen(*, length: int = 32) -> TensorType:
+    """Generate random bytes for symmetric keys or generic randomness.
+
+    API: keygen(length: int = 32) -> key: u8[length]
+
+    Notes:
+    - Frontend defines the type/shape; backend provides randomness.
+    - Raises ValueError when length <= 0.
+    """
+    if length <= 0:
+        raise ValueError("length must be > 0")
+    return TensorType(UINT8, (length,))
+
+
+@_CRYPTO_MOD.op_def()
+def enc(
+    plaintext: MPObject, key: MPObject, algo: str = "aes-ctr"
+) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """Symmetric encryption with algorithm-aware output sizing.
+
+    API: enc(plaintext: u8[N], key: u8[M], *, algo: str = "aes-ctr") -> ciphertext: u8[N + overhead]
+
+    Supported algorithms and overhead:
+    - "aes-ctr": 16 bytes (nonce only, legacy compatibility)
+    - "aes-gcm": 28 bytes (nonce + 16-byte authentication tag)
+    - "sm4-gcm": 28 bytes (nonce + 16-byte authentication tag)
+
+    The algo parameter is stored in the PFunction attributes for backend use.
+    """
+    pt_ty = plaintext
+    if pt_ty.dtype != UINT8:
+        raise TypeError("enc expects UINT8 plaintext")
+    if len(pt_ty.shape) != 1:
+        raise TypeError("enc expects 1-D plaintext")
+
+    # Validate and get overhead for the specified algorithm
+    overhead = _get_algo_overhead(algo)
+    length = pt_ty.shape[0]
+    if length >= 0 and overhead >= 0:
+        outs_info = (TensorType(UINT8, (length + overhead,)),)
+    else:
+        # Unknown length or overhead, return dynamic length
+        outs_info = (TensorType(UINT8, (-1,)),)
+
+    ins_info = (TensorType.from_obj(pt_ty), TensorType.from_obj(key))
+    pfunc = PFunction(
+        fn_type="crypto.enc",
+        ins_info=ins_info,
+        outs_info=outs_info,
+        algo=algo,
+    )
+    _, treedef = tree_flatten(outs_info[0])
+    return pfunc, [plaintext, key], treedef
+
+
+@_CRYPTO_MOD.op_def()
+def dec(
+    ciphertext: MPObject, key: MPObject, algo: str = "aes-ctr"
+) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """Symmetric decryption with algorithm-aware input sizing.
+
+    API: dec(ciphertext: u8[N + overhead], key: u8[M], *, algo: str = "aes-ctr") -> plaintext: u8[N]
+
+    Supported algorithms and overhead:
+    - "aes-ctr": 16 bytes (nonce only, legacy compatibility)
+    - "aes-gcm": 28 bytes (nonce + 16-byte authentication tag)
+    - "sm4-gcm": 28 bytes (nonce + 16-byte authentication tag)
+
+    The algo parameter is stored in the PFunction attributes for backend use.
+    Backend is responsible for parsing the ciphertext format according to algo.
+    """
+    ct_ty = ciphertext
+    if ct_ty.dtype != UINT8:
+        raise TypeError("dec expects UINT8 ciphertext")
+    if len(ct_ty.shape) != 1:
+        raise TypeError("dec expects 1-D ciphertext")
+
+    # Validate and get overhead for the specified algorithm
+    overhead = _get_algo_overhead(algo)
+    length = ct_ty.shape[0]
+
+    # Validate minimum ciphertext length
+    if length >= 0 and overhead >= 0 and length < overhead:
+        raise TypeError(
+            f"dec expects ciphertext with at least {overhead} bytes for algo='{algo}', but got {length} bytes"
+        )
+
+    # Compute output plaintext length
+    if length >= 0 and overhead >= 0:
+        outs_info = (TensorType(UINT8, (length - overhead,)),)
+    else:
+        # Unknown length or overhead, return dynamic length
+        outs_info = (TensorType(UINT8, (-1,)),)
+
+    ins_info = (TensorType.from_obj(ct_ty), TensorType.from_obj(key))
+    pfunc = PFunction(
+        fn_type="crypto.dec",
+        ins_info=ins_info,
+        outs_info=outs_info,
+        algo=algo,
+    )
+    _, treedef = tree_flatten(outs_info[0])
+    return pfunc, [ciphertext, key], treedef
+
+
+@_CRYPTO_MOD.op_def()
+def kem_keygen(suite: str = "x25519") -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """KEM-style keypair generation: returns (sk, pk) bytes.
+
+    API: kem_keygen(suite: str = "x25519") -> (sk: u8[32], pk: u8[32])
+
+    The suite parameter is stored in the PFunction attributes for backend use.
+    """
+    if suite == "x25519":
+        sk_ty = TensorType(UINT8, (32,))
+        pk_ty = TensorType(UINT8, (32,))
+    else:
+        # Unknown suite, return dynamic lengths
+        sk_ty = TensorType(UINT8, (-1,))
+        pk_ty = TensorType(UINT8, (-1,))
+    outs_info = (sk_ty, pk_ty)
+
+    pfunc = PFunction(
+        fn_type="crypto.kem_keygen",
+        ins_info=(),
+        outs_info=outs_info,
+        suite=suite,
+    )
+    _, treedef = tree_flatten(outs_info)
+    return pfunc, [], treedef
+
+
+@_CRYPTO_MOD.op_def()
+def kem_derive(
+    sk: MPObject, peer_pk: MPObject, suite: str = "x25519"
+) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """KEM-style shared secret derivation: returns secret bytes.
+
+    API: kem_derive(sk: u8[32], peer_pk: u8[32], suite: str = "x25519") -> secret: u8[32]
+
+    The suite parameter is stored in the PFunction attributes for backend use.
+    """
+    # Validate input types
+    if sk.dtype != UINT8:
+        raise TypeError("kem_derive expects UINT8 secret key")
+    if peer_pk.dtype != UINT8:
+        raise TypeError("kem_derive expects UINT8 peer public key")
+    if len(sk.shape) != 1 or len(peer_pk.shape) != 1:
+        raise TypeError("kem_derive expects 1-D inputs")
+
+    if suite == "x25519":
+        if sk.shape[0] != 32 or peer_pk.shape[0] != 32:
+            raise TypeError("kem_derive expects 32-byte keys for suite 'x25519'")
+        secret_ty = TensorType(UINT8, (32,))
+    else:
+        # Unknown suite, return dynamic length
+        secret_ty = TensorType(UINT8, (-1,))
+    outs_info = (secret_ty,)
+
+    ins_info = (TensorType.from_obj(sk), TensorType.from_obj(peer_pk))
+    pfunc = PFunction(
+        fn_type="crypto.kem_derive",
+        ins_info=ins_info,
+        outs_info=outs_info,
+        suite=suite,
+    )
+    _, treedef = tree_flatten(outs_info[0])
+    return pfunc, [sk, peer_pk], treedef
+
+
+@_CRYPTO_MOD.op_def()
+def hkdf(
+    secret: MPObject, info: str, hash: str = "SHA-256"
+) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+    """HKDF-style key derivation: returns a 32-byte key.
+
+    API: hkdf(secret: u8[N], info: str, hash: str = "SHA-256") -> key: u8[32]
+
+    The hash parameter is stored in the PFunction attributes for backend use.
+    """
+    # Validate input types
+    if secret.dtype != UINT8:
+        raise TypeError("hkdf expects UINT8 secret")
+    if len(secret.shape) != 1:
+        raise TypeError("hkdf expects 1-D secret")
+
+    if hash == "SHA-256" or hash == "SM3":
+        outs_info = (TensorType(UINT8, (32,)),)
+    else:
+        # Unknown hash, return dynamic length
+        outs_info = (TensorType(UINT8, (-1,)),)
+
+    ins_info = (TensorType.from_obj(secret),)
+    pfunc = PFunction(
+        fn_type="crypto.hkdf",
+        ins_info=ins_info,
+        outs_info=outs_info,
+        hash=hash,
+        info=info,
+    )
+    _, treedef = tree_flatten(outs_info[0])
+    return pfunc, [secret], treedef

mplang/{ops → v1/ops}/fhe.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from mplang.core import UINT8, TensorType
-from mplang.ops.base import stateless_mod
+from mplang.v1.core import UINT8, TensorType
+from mplang.v1.ops.base import stateless_mod
 
 _fhe_MOD = stateless_mod("fhe")
 

mplang/{ops → v1/ops}/jax_cc.py

@@ -14,16 +14,18 @@
 
 from __future__ import annotations
 
+import logging
 from collections.abc import Callable
 from typing import Any
 
 import jax
 import jax.numpy as jnp
+from jax import export
 from jax.tree_util import PyTreeDef, tree_flatten
 
-from mplang.core import MPObject, PFunction, TensorType, get_fn_name
-from mplang.ops.base import FeOperation, stateless_mod
-from mplang.utils.func_utils import normalize_fn
+from mplang.v1.core import MPObject, PFunction, TensorType, get_fn_name
+from mplang.v1.ops.base import FeOperation, stateless_mod
+from mplang.v1.utils.func_utils import normalize_fn
 
 # Enable 64-bit precision for JAX to match tensor types
 jax.config.update("jax_enable_x64", True)
@@ -36,7 +38,8 @@ def jax2stablehlo(
 
     Translates high-level JAX functions into StableHLO MLIR representations,
     enabling execution on JAX backends across different processes and platforms.
-    Uses the standard JAX compilation pipeline: jit → trace → lower → StableHLO MLIR.
+    Uses a hybrid approach: traditional JAX trace/lower for compilation compatibility,
+    with stable jax.export API for parameter tracking.
 
     Args:
         is_variable: Predicate function to classify parameters as variables vs. constants.
@@ -52,34 +55,6 @@ def jax2stablehlo(
                    Non-variable parameters are captured as compile-time constants within
                    the PFunction body, while variables become runtime input parameters.
             - PyTreeDef: Tree structure template for reconstructing nested output values
-
-    Rationale:
-        JAX Serialization Options Analysis:
-        1. jax.export (JAX ≥0.4.35) - Official export API with StableHLO backend
-        2. HLO protobuf - Raw XLA HloModule serialization
-        3. HLO text - Human-readable HLO representation
-        4. StableHLO MLIR - Portable intermediate representation
-        5. JAX compiled object pickling - Limited to same-process execution
-
-        Current Choice: StableHLO MLIR
-        Advantages:
-        - ✅ Available in current JAX version (0.4.34)
-        - ✅ Cross-version compatibility guaranteed by StableHLO design
-        - ✅ Direct compilation support via XLA client.compile(mlir_string)
-        - ✅ Handles complex functions (multi-input/output, control flow)
-        - ✅ Preserves numerical precision
-        - ✅ Platform-independent representation
-
-        Alternative Options Issues:
-        - jax.export: Not available in JAX 0.4.34
-        - HLO protobuf: Version compatibility issues with StableHLO parser
-        - HLO text: Parser compatibility issues with XLA client
-        - Pickle: Cannot serialize XLA LoadedExecutable objects
-
-        Future Migration Path:
-        - JAX ≥0.4.35: Migrate to jax.export.export() + jax.export.deserialize()
-        - JAX ≥0.5.x: Consider new portable formats if available
-        - Long-term: Adopt official JAX serialization standards as they mature
     """
     # Flatten (args, kwargs) and capture immediates using the moved logic from primitive.py
     normalized_fn, in_vars = normalize_fn(flat_fn, args, kwargs, is_variable)
@@ -89,47 +64,39 @@ def jax2stablehlo(
         jax.ShapeDtypeStruct(arg.shape, jnp.dtype(arg.dtype.name)) for arg in in_vars
     ]
 
-    # Standard JAX serialization pipeline: jit → trace → lower → StableHLO MLIR
+    # Hybrid approach: Use standard JAX trace/lower for compatibility, but jax.export for parameter tracking
     jitted_fn = jax.jit(normalized_fn)
     traced = jitted_fn.trace(jax_params)
     lowered = traced.lower()
 
-    # Get StableHLO MLIR representation - the portable format
-    # compiler_ir("stablehlo") returns jaxlib.mlir.ir.Module object
-    # str() converts to serializable text format
+    # Get StableHLO MLIR representation using traditional approach
     stablehlo_mlir = lowered.compiler_ir("stablehlo")
     mlir_text = str(stablehlo_mlir)
 
-    # Get output info and tree structure for result reconstruction after remote execution
+    # Get output info using traditional approach
     out_info_flat, out_tree = tree_flatten(lowered.out_info)
     out_info_flat = [TensorType.from_obj(info) for info in out_info_flat]
 
-    # Extract argument keep mapping to handle JAX's unused parameter elimination
-    # JAX can eliminate unused parameters during compilation, but the runtime still
-    # receives all original arguments. We need the mapping to filter them correctly.
+    # Extract argument keep mapping using stable jax.export API for parameter tracking
+    # We use jax.export only for getting the kept_var_idx information, not for the main compilation
     arg_keep_map = None
     original_arg_count = len(in_vars)
 
     try:
-        # Access JAX internal kept_var_idx - the authoritative source
-        # This tells us exactly which original parameters survived compilation
-        compile_args = lowered._lowering.compile_args
-        kept_var_idx = compile_args["kept_var_idx"]
-
-        kept_indices = sorted(kept_var_idx)
-        if len(kept_indices) < original_arg_count:
-            arg_keep_map = kept_indices
-
-    except (AttributeError, KeyError, TypeError) as e:
-        # JAX internal API is not available or changed
-        # This is a hard error - we cannot reliably handle unused parameters
-        # without knowing exactly which ones were kept
-        raise RuntimeError(
-            f"Cannot access JAX's kept_var_idx to handle unused parameter elimination. "
-            f"This function may have unused parameters that JAX optimized away, "
-            f"but we cannot determine which ones without the internal API. "
-            f"Original error: {e}"
-        ) from e
+        # Use jax.export just to get the stable parameter tracking information
+        export_fn = export.export(jitted_fn)
+        exported = export_fn(jax_params)
+        kept_var_idx = exported.module_kept_var_idx
+        if kept_var_idx is not None and len(kept_var_idx) < original_arg_count:
+            # JAX eliminated some unused parameters during compilation
+            # Keep the indices in sorted order for consistent mapping
+            arg_keep_map = sorted(kept_var_idx)
+    except Exception as e:
+        # Fallback: if jax.export fails, we can still use the compiled result without parameter tracking
+        # This ensures backward compatibility even if export has issues
+        logging.warning(
+            f"jax.export failed to get kept_var_idx, proceeding without it. Error: {e}"
+        )
 
     # This format tells JaxRT how to handle the compiled result
     pfn_kwargs: dict[str, Any] = {

mplang/v1/ops/nnx_cc.py

@@ -0,0 +1,168 @@
+# Copyright 2025 Ant Group Co., 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 __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jax import export
+from jax.tree_util import PyTreeDef, tree_flatten
+
+from mplang.v1.core import MPObject, PFunction, TensorType, get_fn_name
+from mplang.v1.ops.base import FeOperation, stateless_mod
+from mplang.v1.utils.func_utils import normalize_fn
+
+# Enable 64-bit precision for JAX to match tensor types
+jax.config.update("jax_enable_x64", True)
+
+
+def nnx2stablehlo(
+    is_variable: Callable[[Any], bool], flat_fn: Any, *args: Any, **kwargs: Any
+) -> tuple[PFunction, list[Any], PyTreeDef]:
+    """Compile NNX function to StableHLO MLIR format for remote execution.
+
+    Translates high-level NNX functions into StableHLO MLIR representations,
+    enabling execution on JAX backends across different processes and platforms.
+    Uses a hybrid approach: traditional NNX trace/lower for compilation compatibility,
+    with stable jax.export API for parameter tracking.
+
+    Args:
+        is_variable: Predicate function to classify parameters as variables vs. constants.
+                    Returns True for parameters that should be treated as PFunction inputs.
+        flat_fn: NNX function to be compiled into StableHLO format
+        *args: Positional arguments passed to the function during compilation
+        **kwargs: Keyword arguments passed to the function during compilation
+
+    Returns:
+        tuple[PFunction, list, PyTreeDef]: Compilation artifacts containing:
+            - PFunction: Serialized function with embedded MLIR text and type metadata
+            - list: Extracted variable parameters (those satisfying is_variable predicate).
+                   Non-variable parameters are captured as compile-time constants within
+                   the PFunction body, while variables become runtime input parameters.
+            - PyTreeDef: Tree structure template for reconstructing nested output values
+    """
+    # Flatten (args, kwargs) and capture immediates using the moved logic from primitive.py
+    normalized_fn, in_vars = normalize_fn(flat_fn, args, kwargs, is_variable)
+
+    # Convert TensorType in_vars to ShapeDtypeStruct for JAX tracing
+    jax_params = [
+        jax.ShapeDtypeStruct(arg.shape, jnp.dtype(arg.dtype.name)) for arg in in_vars
+    ]
+
+    # NNX compilation pipeline using JAX export API: nnx.jit → jax.export → StableHLO MLIR
+    # Use nnx.jit for NNX-specific functionality, then jax.export for stable parameter handling
+    nnx_jitted = nnx.jit(normalized_fn)
+
+    # Extract the underlying JAX function for jax.export compatibility
+    # nnx.jit wraps a JAX function, and we can access it via .fun attribute
+    underlying_jax_fn = nnx_jitted.fun
+
+    # Hybrid approach: Use NNX trace/lower for compilation, but jax.export for parameter tracking
+    # Use traditional nnx.jit → trace → lower for compatibility with argument structure
+    nnx_traced = nnx_jitted.trace(jax_params)
+    nnx_lowered = nnx_traced.lower()
+
+    # Get StableHLO MLIR representation using traditional NNX approach
+    # NNX lowered object wraps JAX lowered, so we access the inner JAX lowered object
+    jax_lowered = nnx_lowered.lowered
+    stablehlo_mlir = jax_lowered.compiler_ir("stablehlo")
+    mlir_text = str(stablehlo_mlir)
+
+    # Get output info using traditional NNX approach
+    # NNX captures output in (args, kwargs, result) format, so we need to extract just the result part
+    raw_out_info = jax_lowered.out_info
+    if isinstance(raw_out_info, tuple) and len(raw_out_info) == 3:
+        # NNX format: (args, kwargs, result) - extract just the result
+        _, _, actual_out_info = raw_out_info
+        out_info_flat, out_tree = tree_flatten(actual_out_info)
+    else:
+        # Fallback to direct format (shouldn't happen with NNX, but just in case)
+        out_info_flat, out_tree = tree_flatten(raw_out_info)
+
+    out_info_flat = [TensorType.from_obj(info) for info in out_info_flat]
+
+    # Extract argument keep mapping using stable jax.export API for parameter tracking
+    # We use the underlying JAX function with jax.export only for parameter tracking
+    arg_keep_map = None
+    original_arg_count = len(in_vars)
+
+    try:
+        # Use jax.export with the underlying JAX function just to get stable parameter tracking
+        export_fn = export.export(jax.jit(underlying_jax_fn))
+        exported = export_fn(jax_params)
+        kept_var_idx = exported.module_kept_var_idx
+        if kept_var_idx is not None and len(kept_var_idx) < original_arg_count:
+            # JAX eliminated some unused parameters during compilation
+            # Keep the indices in sorted order for consistent mapping
+            arg_keep_map = sorted(kept_var_idx)
+    except Exception as e:
+        # Fallback: if jax.export fails, we can still use the compiled result without parameter tracking
+        # This ensures backward compatibility even if export has issues
+        logging.warning(
+            f"jax.export failed to get kept_var_idx, proceeding without it. Error: {e}"
+        )
+
+    # This format tells JaxRT how to handle the compiled result
+    # Use the same format as JAX since NNX compiles to the same backend
+    pfn_kwargs: dict[str, Any] = {
+        "fn_type": "mlir.stablehlo",  # Key: specify StableHLO MLIR format
+        "ins_info": tuple(TensorType.from_obj(x) for x in in_vars),
+        "outs_info": tuple(out_info_flat),
+        "fn_name": get_fn_name(flat_fn),
+        "fn_text": mlir_text,  # MLIR text, serializable for transmission
+    }
+
+    if arg_keep_map is not None:
+        pfn_kwargs["arg_keep_map"] = arg_keep_map
+
+    pfn = PFunction(**pfn_kwargs)
+    return pfn, in_vars, out_tree
+
+
+class NnxRunner(FeOperation):
+    """NNX function runner frontend operation."""
+
+    def trace(
+        self, nnx_fn: Callable, *args: Any, **kwargs: Any
+    ) -> tuple[PFunction, list[MPObject], PyTreeDef]:
+        """
+        NNX compilation helper function.
+
+        Compiles an NNX function to StableHLO format and returns the PFunction
+        along with variable arguments for evaluation.
+
+        Args:
+            nnx_fn: The NNX function to compile
+            *args: Positional arguments to the function
+            **kwargs: Keyword arguments to the function
+
+        Returns:
+            tuple[PFunction, list[MPObject], PyTreeDef]: The compiled PFunction, input variables, and output tree
+        """
+
+        def is_variable(arg: Any) -> bool:
+            return isinstance(arg, MPObject)
+
+        pfunc, in_vars, out_tree = nnx2stablehlo(is_variable, nnx_fn, *args, **kwargs)
+        return pfunc, in_vars, out_tree
+
+
+_NNX_MOD = stateless_mod("nnx")
+
+run_nnx = NnxRunner(_NNX_MOD, "run")

mplang/{ops → v1/ops}/phe.py

@@ -14,21 +14,34 @@
 
 """PHE (Partially Homomorphic Encryption) frontend operations."""
 
-from mplang.core import UINT8, TensorType
-from mplang.ops.base import stateless_mod
+from mplang.v1.core import UINT8, TensorType
+from mplang.v1.ops.base import stateless_mod
 
 _PHE_MOD = stateless_mod("phe")
 
 
 @_PHE_MOD.simple_op()
 def keygen(
-    *, scheme: str = "paillier", key_size: int = 2048
+    *,
+    scheme: str = "paillier",
+    key_size: int = 2048,
+    max_value: int | None = None,
+    fxp_bits: int | None = None,
 ) -> tuple[TensorType, TensorType]:
     """Generate a PHE key pair: returns (public_key, private_key).
 
     Keys are represented with a sentinel TensorType UINT8[(-1, 0)] to indicate
     non-structural, backend-only handles. Runtime validation will treat this
     shape as an opaque placeholder and skip dtype/shape checks.
+
+    Attributes (forwarded to backend):
+        scheme: PHE scheme (default: 'paillier')
+        key_size: Modulus size in bits (default: 2048)
+        max_value: Optional range-encoding bound B. If provided, the backend will
+            encode/decode integers/floats within [-B, B] and treat (B, N-B) as overflow.
+            Pick B to exceed the largest intermediate magnitude you expect in homomorphic
+            combinations. If omitted, backend default is used (currently 2**32).
+        fxp_bits: Optional fixed-point fractional bits for float encoding (default backend value).
     """
     key_spec = TensorType(UINT8, (-1, 0))
     return key_spec, key_spec

mplang/{ops → v1/ops}/spu.py

@@ -23,9 +23,9 @@ import spu.utils.frontend as spu_fe
 from jax import ShapeDtypeStruct
 from jax.tree_util import PyTreeDef, tree_flatten
 
-from mplang.core import MPObject, PFunction, TensorType, get_fn_name
-from mplang.ops.base import stateless_mod
-from mplang.utils.func_utils import normalize_fn
+from mplang.v1.core import MPObject, PFunction, TensorType, get_fn_name
+from mplang.v1.ops.base import stateless_mod
+from mplang.v1.utils.func_utils import normalize_fn
 
 
 class Visibility: