mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev271__py3-none-any.whl

mplang/v1/ops/fhe.py

@@ -1,272 +0,0 @@
-# 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 mplang.v1.core import UINT8, TensorType
-from mplang.v1.ops.base import stateless_mod
-
-_fhe_MOD = stateless_mod("fhe")
-
-
-@_fhe_MOD.simple_op()
-def keygen(
-    *,
-    scheme: str = "CKKS",
-    poly_modulus_degree: int = 8192,
-    coeff_mod_bit_sizes: tuple[int, ...] | None = None,
-    global_scale: int | None = None,
-    plain_modulus: int | None = None,
-) -> tuple[TensorType, TensorType, TensorType]:
-    """Generate an FHE key pair for Vector backend: returns (private_context, public_context, evaluation_context).
-
-    Args:
-        scheme: FHE scheme to use ("CKKS" for approximate, "BFV" for exact integer)
-        poly_modulus_degree: Polynomial modulus degree (default: 8192)
-        coeff_mod_bit_sizes: Coefficient modulus bit sizes for CKKS (optional)
-        global_scale: Global scale for CKKS (optional)
-        plain_modulus: Plain modulus for BFV (optional)
-
-    Returns:
-        Tuple of (private_context, public_context, evaluation_context) represented as UINT8[(-1, 0)]
-
-    Contexts are represented with a sentinel TensorType UINT8[(-1, 0)] to indicate
-    non-structural, backend-only handles.
-
-    Note: Vector backend only supports 1D data. For multi-dimensional tensors,
-    use mplang.ops.fhe instead.
-    """
-    if scheme not in ("CKKS", "BFV"):
-        raise ValueError("Unsupported scheme. Choose either 'CKKS' or 'BFV'.")
-    if scheme == "CKKS":
-        assert plain_modulus is None, "plain_modulus is not used in CKKS scheme."
-    context_spec = TensorType(UINT8, (-1, 0))
-    return context_spec, context_spec, context_spec
-
-
-@_fhe_MOD.simple_op()
-def encrypt(plaintext: TensorType, context: TensorType) -> TensorType:
-    """Encrypt plaintext using FHE Vector backend: returns ciphertext with same semantic type.
-
-    Args:
-        plaintext: Data to encrypt (scalar or 1D vector only)
-        context: FHE context (private or public)
-
-    Returns:
-        Ciphertext with same semantic type as plaintext
-
-    Raises:
-        ValueError: If plaintext has more than 1 dimension
-
-    Note: Vector backend only supports scalars (shape=()) and 1D vectors (shape=(n,)).
-    For multi-dimensional data, use mplang.ops.fhe.encrypt instead.
-    """
-    _ = context
-    if len(plaintext.shape) > 1:
-        raise ValueError(
-            f"FHE Vector backend only supports 1D data. Got shape {plaintext.shape}. "
-            "Use mplang.ops.fhe for multi-dimensional tensors."
-        )
-    return plaintext
-
-
-@_fhe_MOD.simple_op()
-def decrypt(ciphertext: TensorType, context: TensorType) -> TensorType:
-    """Decrypt ciphertext using FHE Vector backend: returns plaintext with same semantic type.
-
-    Args:
-        ciphertext: Encrypted data to decrypt (scalar or 1D vector)
-        context: FHE context (must be private context with secret key)
-
-    Returns:
-        Plaintext with same semantic type as ciphertext
-
-    Note: Ciphertext encrypted with public context can be decrypted with
-    the corresponding private context.
-    """
-    _ = context
-    return ciphertext
-
-
-@_fhe_MOD.simple_op()
-def add(operand1: TensorType, operand2: TensorType) -> TensorType:
-    """Add two FHE operands (ciphertext + ciphertext or ciphertext + plaintext).
-
-    Args:
-        operand1: First operand (ciphertext or plaintext, scalar or 1D vector)
-        operand2: Second operand (ciphertext or plaintext, scalar or 1D vector)
-
-    Returns:
-        Result of homomorphic addition
-
-    Raises:
-        ValueError: If operands have incompatible shapes or dtypes
-
-    Note: At least one operand must be ciphertext. Both operands must have
-    the same shape (no broadcasting in Vector backend).
-    """
-    assert operand1.dtype == operand2.dtype, (
-        f"Operand dtypes must match, got {operand1.dtype} and {operand2.dtype}."
-    )
-    assert operand1.shape == operand2.shape, (
-        f"Operand shapes must match, got {operand1.shape} and {operand2.shape}."
-    )
-    return operand1
-
-
-@_fhe_MOD.simple_op()
-def sub(operand1: TensorType, operand2: TensorType) -> TensorType:
-    """Subtract two FHE operands (ciphertext - ciphertext or ciphertext - plaintext).
-
-    Args:
-        operand1: First operand (ciphertext or plaintext, scalar or 1D vector)
-        operand2: Second operand (ciphertext or plaintext, scalar or 1D vector)
-
-    Returns:
-        Result of homomorphic subtraction
-
-    Raises:
-        ValueError: If operands have incompatible shapes or dtypes
-
-    Note: At least one operand must be ciphertext. Both operands must have
-    the same shape (no broadcasting in Vector backend).
-    """
-    assert operand1.dtype == operand2.dtype, (
-        f"Operand dtypes must match, got {operand1.dtype} and {operand2.dtype}."
-    )
-    assert operand1.shape == operand2.shape, (
-        f"Operand shapes must match, got {operand1.shape} and {operand2.shape}."
-    )
-    return operand1
-
-
-@_fhe_MOD.simple_op()
-def mul(operand1: TensorType, operand2: TensorType) -> TensorType:
-    """Multiply two FHE operands (ciphertext * ciphertext or ciphertext * plaintext).
-
-    Args:
-        operand1: First operand (ciphertext or plaintext, scalar or 1D vector)
-        operand2: Second operand (ciphertext or plaintext, scalar or 1D vector)
-
-    Returns:
-        Result of homomorphic multiplication
-
-    Raises:
-        ValueError: If operands have incompatible shapes or dtypes
-
-    Note: At least one operand must be ciphertext. Both operands must have
-    the same shape (no broadcasting in Vector backend).
-    For BFV scheme, plaintext operands must be integers.
-    """
-    assert operand1.dtype == operand2.dtype, (
-        f"Operand dtypes must match, got {operand1.dtype} and {operand2.dtype}."
-    )
-    assert operand1.shape == operand2.shape, (
-        f"Operand shapes must match, got {operand1.shape} and {operand2.shape}."
-    )
-    return operand1
-
-
-@_fhe_MOD.simple_op()
-def dot(operand1: TensorType, operand2: TensorType) -> TensorType:
-    """Compute dot product of FHE operands (ciphertext · ciphertext or ciphertext · plaintext).
-
-    Args:
-        operand1: First operand (ciphertext or plaintext, must be 1D vector)
-        operand2: Second operand (ciphertext or plaintext, must be 1D vector)
-
-    Returns:
-        Scalar result of homomorphic dot product (shape=())
-
-    Raises:
-        ValueError: If operands are not 1D vectors or have different lengths
-
-    Note: Both operands must be 1D vectors (not scalars). For scalar multiplication,
-    use mul() instead. This operation always returns a scalar.
-    """
-    if len(operand1.shape) != 1:
-        raise ValueError(
-            f"Dot product requires 1D vectors, got shape {operand1.shape} for operand1"
-        )
-    if len(operand2.shape) != 1:
-        raise ValueError(
-            f"Dot product requires 1D vectors, got shape {operand2.shape} for operand2"
-        )
-    if operand1.shape[0] != operand2.shape[0]:
-        raise ValueError(
-            f"Dot product dimension mismatch: {operand1.shape[0]} vs {operand2.shape[0]}"
-        )
-
-    # Dot product of 1D vectors returns a scalar
-    return TensorType(operand1.dtype, ())
-
-
-@_fhe_MOD.simple_op()
-def polyval(ciphertext: TensorType, coeffs: TensorType) -> TensorType:
-    """Evaluate polynomial on encrypted data with plaintext coefficients.
-
-    Args:
-        ciphertext: Encrypted data (scalar or 1D vector)
-        coeffs: Plaintext polynomial coefficients as 1D array [c0, c1, c2, ...]
-                representing c0 + c1*x + c2*x^2 + ...
-
-    Returns:
-        Result of polynomial evaluation with same shape and dtype as ciphertext
-
-    Raises:
-        ValueError: If coefficients array is not 1D or has fewer than 2 elements
-
-    Note: Polynomial must have degree >= 1 (at least 2 coefficients required).
-    Constant polynomials (degree 0, single coefficient) are NOT supported due to
-    TenSEAL limitation. For constant values, use: ct * 0 + constant instead.
-    For BFV scheme, coefficients must be integers.
-
-    Common use case - Sigmoid approximation:
-        sigmoid_coeffs = [0.5, 0.15012, 0.0, -0.0018027]
-        result = polyval(ciphertext, sigmoid_coeffs)
-    """
-    if len(coeffs.shape) != 1:
-        raise ValueError(
-            f"Polynomial coefficients must be 1D array, got shape {coeffs.shape}"
-        )
-    _ = coeffs
-    return ciphertext
-
-
-@_fhe_MOD.simple_op()
-def negate(ciphertext: TensorType) -> TensorType:
-    """Negate encrypted data (unary minus).
-
-    Args:
-        ciphertext: Encrypted data (scalar or 1D vector)
-
-    Returns:
-        Negated ciphertext with same shape and dtype
-
-    Note: Equivalent to multiplying by -1.
-    """
-    return ciphertext
-
-
-@_fhe_MOD.simple_op()
-def square(ciphertext: TensorType) -> TensorType:
-    """Square encrypted data (element-wise).
-
-    Args:
-        ciphertext: Encrypted data (scalar or 1D vector)
-
-    Returns:
-        Squared ciphertext with same shape and dtype
-
-    Note: More efficient than mul(ciphertext, ciphertext) in some FHE schemes.
-    """
-    return ciphertext

mplang/v1/ops/jax_cc.py

@@ -1,147 +0,0 @@
-# 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 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 jax2stablehlo(
-    is_variable: Callable[[Any], bool], flat_fn: Any, *args: Any, **kwargs: Any
-) -> tuple[PFunction, list, PyTreeDef]:
-    """Compile JAX function to StableHLO MLIR format for remote execution.
-
-    Translates high-level JAX functions into StableHLO MLIR representations,
-    enabling execution on JAX backends across different processes and platforms.
-    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.
-                    Returns True for parameters that should be treated as PFunction inputs.
-        flat_fn: JAX 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
-    ]
-
-    # 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 using traditional approach
-    stablehlo_mlir = lowered.compiler_ir("stablehlo")
-    mlir_text = str(stablehlo_mlir)
-
-    # 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 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:
-        # 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] = {
-        "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 JaxRunner(FeOperation):
-    """JAX function runner frontend operation."""
-
-    def trace(
-        self, jax_fn: Callable, *args: Any, **kwargs: Any
-    ) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-        """
-        JAX compilation helper function.
-
-        Compiles a JAX function to StableHLO format and returns the PFunction
-        along with variable arguments for evaluation.
-
-        Args:
-            jax_fn: The JAX 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 = jax2stablehlo(is_variable, jax_fn, *args, **kwargs)
-        return pfunc, in_vars, out_tree
-
-
-_JAX_MOD = stateless_mod("jax")
-
-run_jax = JaxRunner(_JAX_MOD, "run")

mplang/v1/ops/nnx_cc.py

@@ -1,168 +0,0 @@
-# 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")