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

mplang/{v2/dialects → dialects}/tee.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""TEE (Trusted Execution Environment) dialect for mplang.v2 EDSL.
+"""TEE (Trusted Execution Environment) dialect for mplang EDSL.
 
 This dialect provides primitives for TEE remote attestation, enabling secure
 computation where:
@@ -38,7 +38,7 @@ Supported Platforms:
 
 Example:
 ```python
-from mplang.v2.dialects import tee, crypto
+from mplang.dialects import tee, crypto
 
 # On TEE side: generate keypair and quote
 sk, pk = crypto.kem_keygen("x25519")
@@ -62,10 +62,10 @@ from __future__ import annotations
 
 from typing import Any, ClassVar, Literal
 
-import mplang.v2.edsl as el
-import mplang.v2.edsl.typing as elt
-from mplang.v2.dialects.crypto import PublicKeyType
-from mplang.v2.edsl import serde
+import mplang.edsl as el
+import mplang.edsl.typing as elt
+from mplang.dialects.crypto import PublicKeyType
+from mplang.edsl import serde
 
 # ==============================================================================
 # --- Type Definitions

mplang/{v2/dialects → dialects}/tensor.py

@@ -58,10 +58,10 @@ import numpy as np
 from jax import ShapeDtypeStruct
 from jax.tree_util import PyTreeDef, tree_flatten
 
-import mplang.v2.edsl as el
-import mplang.v2.edsl.typing as elt
-from mplang.v1.utils.func_utils import normalize_fn
-from mplang.v2.dialects import dtypes
+import mplang.edsl as el
+import mplang.edsl.typing as elt
+from mplang.dialects import dtypes
+from mplang.utils import normalize_fn
 
 run_jax_p = el.Primitive[Any]("tensor.run_jax")
 constant_p = el.Primitive[el.Object]("tensor.constant")
@@ -77,7 +77,7 @@ class RunJaxCompilation:
 
     fn: Callable[..., Any]
     stablehlo: str
-    out_tree: PyTreeDef
+    out_tree: PyTreeDef  # type: ignore
     output_types: list[elt.BaseType]
     arg_keep_map: list[int] | None = None
 

mplang/{v2/edsl → edsl}/__init__.py

@@ -17,8 +17,8 @@
 This module keeps the surface area intentionally small so downstream code can
 simply write::
 
-    import mplang.v2.edsl as el
-    import mplang.v2.edsl.typing as elt
+    import mplang.edsl as el
+    import mplang.edsl.typing as elt
 
 The `el` namespace re-exports the commonly used building blocks (context,
 graph, tracer, primitives, etc.), while the full type system lives under
@@ -27,7 +27,7 @@ graph, tracer, primitives, etc.), while the full type system lives under
 
 from __future__ import annotations
 
-# Re-export the typing module so callers can `import mplang.v2.edsl.typing as elt`
+# Re-export the typing module so callers can `import mplang.edsl.typing as elt`
 from . import typing as typing
 
 # Context management

mplang/{v2/edsl → edsl}/context.py

@@ -21,7 +21,7 @@ This module defines the Context hierarchy:
 
 Contexts can be used directly with Python's 'with' statement:
 
-    from mplang.v2.edsl import Tracer
+    from mplang.edsl import Tracer
 
     tracer = Tracer()
     with tracer:
@@ -46,9 +46,9 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, Self
 
 if TYPE_CHECKING:
-    from mplang.v2.edsl.graph import Graph
-    from mplang.v2.edsl.object import Object
-    from mplang.v2.edsl.primitive import Primitive
+    from mplang.edsl.graph import Graph
+    from mplang.edsl.object import Object
+    from mplang.edsl.primitive import Primitive
 
 
 class Context(ABC):
@@ -257,7 +257,7 @@ def is_tracing() -> bool:
     Returns:
         True if the top of the context stack is a Tracer.
     """
-    from mplang.v2.edsl.tracer import Tracer
+    from mplang.edsl.tracer import Tracer
 
     return isinstance(get_current_context(), Tracer)
 
@@ -280,7 +280,7 @@ def get_default_context() -> Context:
         if _default_context_factory is None:
             raise RuntimeError(
                 "No default context factory registered. "
-                "Ensure mplang.v2.edsl is imported or register a factory manually."
+                "Ensure mplang.edsl is imported or register a factory manually."
             )
         _default_context = _default_context_factory()
     return _default_context

mplang/{v2/edsl → edsl}/graph.py

@@ -27,8 +27,8 @@ Key Design Principles:
 
 Example:
 --------
-    from mplang.v2.edsl.graph import Graph, Operation, Value
-    from mplang.v2.edsl.typing import Tensor, f32
+    from mplang.edsl.graph import Graph, Operation, Value
+    from mplang.edsl.typing import Tensor, f32
 
     graph = Graph()
 
@@ -61,8 +61,8 @@ from collections.abc import Sequence
 from dataclasses import dataclass, field
 from typing import Any, ClassVar
 
-from mplang.v2.edsl import serde
-from mplang.v2.edsl.typing import BaseType
+from mplang.edsl import serde
+from mplang.edsl.typing import BaseType
 
 
 @dataclass
@@ -74,7 +74,7 @@ class Value:
 
     Attributes:
         name: Unique SSA name (e.g., "%0", "%1", ...)
-        type: Type of this value (from mplang.v2.edsl.typing)
+        type: Type of this value (from mplang.edsl.typing)
         defining_op: Operation that produces this value (None for inputs)
         uses: List of operations that consume this value
     """

mplang/{v2/edsl → edsl}/jit.py

@@ -19,12 +19,12 @@ from typing import Any
 
 from jax.tree_util import tree_map
 
-from mplang.v2.edsl.context import (
+from mplang.edsl.context import (
     AbstractInterpreter,
     get_current_context,
     get_default_context,
 )
-from mplang.v2.edsl.tracer import Tracer
+from mplang.edsl.tracer import Tracer
 
 
 def jit(fn: Callable) -> Callable:

mplang/{v2/edsl → edsl}/object.py

@@ -25,7 +25,7 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import Generic, TypeVar
 
-from mplang.v2.edsl.typing import BaseType
+from mplang.edsl.typing import BaseType
 
 T = TypeVar("T", bound=BaseType)
 

mplang/{v2/edsl → edsl}/primitive.py

@@ -27,11 +27,11 @@ from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
 
 from jax.tree_util import tree_map
 
-from mplang.v2.edsl.context import get_current_context, get_default_context
-from mplang.v2.edsl.object import Object
+from mplang.edsl.context import get_current_context, get_default_context
+from mplang.edsl.object import Object
 
 if TYPE_CHECKING:
-    from mplang.v2.edsl.typing import BaseType
+    from mplang.edsl.typing import BaseType
 
 T_Ret = TypeVar("T_Ret")
 
@@ -54,7 +54,7 @@ class Primitive(Generic[T_Ret]):
         >>>
         >>> @encrypt_p.def_abstract_eval
         >>> def encrypt_abstract(x_type):
-        >>>     from mplang.v2.edsl.typing import Vector
+        >>>     from mplang.edsl.typing import Vector
         >>>     return Vector[x_type.dtype, x_type.shape]
         >>>
         >>> # Execution happens via Graph IR → Backend
@@ -91,7 +91,7 @@ class Primitive(Generic[T_Ret]):
         """
         self._impl = fn
         # Register with the global interpreter registry
-        from mplang.v2.edsl.registry import register_impl
+        from mplang.edsl.registry import register_impl
 
         register_impl(self.name, fn)
         return fn

mplang/{v2/edsl → edsl}/printer.py

@@ -18,7 +18,7 @@ from __future__ import annotations
 
 from typing import Any
 
-from mplang.v2.edsl.graph import Graph, Operation, Value
+from mplang.edsl.graph import Graph, Operation, Value
 
 
 class GraphPrinter:

mplang/{v2/edsl → edsl}/serde.py

@@ -20,7 +20,7 @@ cloudpickle. Each type is responsible for its own serialization via the
 `@register_class` decorator pattern.
 
 Usage:
-    from mplang.v2.edsl import serde
+    from mplang.edsl import serde
 
     @serde.register_class
     class MyType:

mplang/{v2/edsl → edsl}/tracer.py

@@ -31,14 +31,14 @@ from typing import TYPE_CHECKING, Any, cast
 
 from jax.tree_util import PyTreeDef, tree_flatten, tree_map
 
-from mplang.v2.edsl.context import Context
-from mplang.v2.edsl.graph import Graph
-from mplang.v2.edsl.graph import Value as GraphValue
-from mplang.v2.edsl.object import Object
-from mplang.v2.edsl.typing import BaseType
+from mplang.edsl.context import Context
+from mplang.edsl.graph import Graph
+from mplang.edsl.graph import Value as GraphValue
+from mplang.edsl.object import Object
+from mplang.edsl.typing import BaseType
 
 if TYPE_CHECKING:
-    from mplang.v2.edsl.primitive import Primitive
+    from mplang.edsl.primitive import Primitive
 
 
 class TraceObject(Object):
@@ -48,7 +48,7 @@ class TraceObject(Object):
     All operations delegate to primitives which record into Graph.
 
     Example:
-        >>> from mplang.v2.edsl import trace
+        >>> from mplang.edsl import trace
         >>> def compute(x, y):
         ...     z = x + y  # TraceObject.__add__ → add_p.bind(x, y)
         ...     return z

mplang/{v2/edsl → edsl}/typing.py

@@ -124,7 +124,7 @@ from __future__ import annotations
 
 from typing import Any, ClassVar, Generic, TypeVar
 
-from mplang.v2.edsl import serde
+from mplang.edsl import serde
 
 # ==============================================================================
 # --- Base Type & Type Aliases

mplang/{v2/kernels → kernels}/ldpc.cpp

@@ -27,13 +27,13 @@ extern "C" {
 
 /**
  * @brief LDPC Encoding: Compute Syndrome s = H * x
- * 
+ *
  * H is a sparse M x N binary matrix (CSR format).
  * x is a dense N-vector of 128-bit blocks (N * 16 bytes).
  * s is a dense M-vector of 128-bit blocks (M * 16 bytes).
- * 
+ *
  * Logic: For each row i of H, s[i] = XOR(x[j]) for all j where H[i, j] = 1.
- * 
+ *
  * @param message_ptr  Pointer to message x (N * 2 uint64_t)
  * @param indices_ptr  Pointer to CSR indices (uint64_t)
  * @param indptr_ptr   Pointer to CSR indptr (M+1 uint64_t)
@@ -41,17 +41,17 @@ extern "C" {
  * @param m            Number of rows in H (syndrome length)
  * @param n            Number of cols in H (message length)
  */
-void ldpc_encode(const uint64_t* message_ptr, 
-                 const uint64_t* indices_ptr, 
-                 const uint64_t* indptr_ptr, 
-                 uint64_t* output_ptr, 
-                 uint64_t m, 
+void ldpc_encode(const uint64_t* message_ptr,
+                 const uint64_t* indices_ptr,
+                 const uint64_t* indptr_ptr,
+                 uint64_t* output_ptr,
+                 uint64_t m,
                  uint64_t n) {
-    
+
     // Check alignment
     // We assume message_ptr and output_ptr are 16-byte aligned for SSE/AVX?
     // JAX/Numpy arrays are usually aligned.
-    
+
     // Cast to __m128i for efficiency
     // But we need to handle potential unaligned access if numpy doesn't align.
     // _mm_loadu_si128 handles unaligned.
@@ -63,10 +63,10 @@ void ldpc_encode(const uint64_t* message_ptr,
     for (uint64_t i = 0; i < m; ++i) {
         // Row i
         __m128i sum = _mm_setzero_si128();
-        
+
         uint64_t start = indptr_ptr[i];
         uint64_t end = indptr_ptr[i+1];
-        
+
         for (uint64_t k = start; k < end; ++k) {
             uint64_t col_idx = indices_ptr[k];
             // XOR accumulation
@@ -74,7 +74,7 @@ void ldpc_encode(const uint64_t* message_ptr,
             __m128i val = _mm_loadu_si128(&x_vec[col_idx]);
             sum = _mm_xor_si128(sum, val);
         }
-        
+
         _mm_storeu_si128(&s_vec[i], sum);
     }
 }

mplang/{v2/kernels → kernels}/okvs.cpp

@@ -86,7 +86,7 @@ extern "C" {
         // 3. Build CSR Structure (Flat Arrays) to replace vector<vector>
         // col_start[j] points to start of column j's rows in flat_rows
         std::vector<int> col_start(m + 1, 0);
-        
+
         // Prefix sum to compute start positions
         // col_start[0] = 0
         // col_start[j+1] = col_start[j] + degree[j]
@@ -96,7 +96,7 @@ extern "C" {
 
         // Total edges = 3 * N implies flat_rows size
         std::vector<int> flat_rows(n * 3);
-        
+
         // Temporary copy of start indices to use as fill pointers
         std::vector<int> fill_ptr = col_start;
 
@@ -106,7 +106,7 @@ extern "C" {
             flat_rows[fill_ptr[r.h2]++] = i;
             flat_rows[fill_ptr[r.h3]++] = i;
         }
-        
+
         // 4. Initialize Peeling
         std::vector<int> peel_stack;
         peel_stack.reserve(m);
@@ -135,7 +135,7 @@ extern "C" {
             int owner_row = -1;
             int start = col_start[j];
             int end = col_start[j+1];
-            
+
             for(int k=start; k<end; ++k) {
                 int r_idx = flat_rows[k];
                 if(!row_removed[r_idx]) {

mplang/{v2/kernels → kernels}/okvs_opt.cpp

@@ -28,7 +28,7 @@
 extern "C" {
 
     // Number of Bins for Mega-Binning strategy.
-    // 1024 bins implies ~1000 items per bin for N=1M, fitting the working set 
+    // 1024 bins implies ~1000 items per bin for N=1M, fitting the working set
     // entirely in L1 cache (32KB/48KB) for maximum performance.
     static const uint64_t NUM_BINS = 1024;
 
@@ -54,14 +54,14 @@ extern "C" {
     inline Indices get_bin_local_indices(uint64_t key, uint64_t m_local, __m128i seed) {
         // Use a distinct seed mix to decorrelate from bin selection
         __m128i k = _mm_set_epi64x(0, key);
-        __m128i s2 = _mm_add_epi64(seed, _mm_set_epi64x(1, 1)); 
+        __m128i s2 = _mm_add_epi64(seed, _mm_set_epi64x(1, 1));
         __m128i h = _mm_aesenc_si128(k, s2);
         h = _mm_aesenc_si128(h, s2);
         h = _mm_aesenc_si128(h, s2);
 
         uint64_t r = _mm_extract_epi64(h, 0);
         Indices idx;
-        
+
         // Fast modulo for local indices
         idx.h1 = r % m_local;
         r = r * 6364136223846793005ULL + 1442695040888963407ULL; // LCG step
@@ -80,10 +80,10 @@ extern "C" {
 
     // Core Peeling Solver for a single Bin
     bool solve_bin(
-        const std::vector<uint64_t>& keys, 
-        const std::vector<__m128i>& vals, 
-        __m128i* P_local, 
-        uint64_t m, 
+        const std::vector<uint64_t>& keys,
+        const std::vector<__m128i>& vals,
+        __m128i* P_local,
+        uint64_t m,
         __m128i seed
     ) {
         uint64_t n = keys.size();
@@ -95,7 +95,7 @@ extern "C" {
         };
         std::vector<Edge> edges(n);
         std::vector<int> col_degree(m, 0);
-        
+
         // 1. Build Local Graph
         for(uint64_t i=0; i<n; ++i) {
              Indices idx = get_bin_local_indices(keys[i], m, seed);
@@ -127,14 +127,14 @@ extern "C" {
 
         std::vector<bool> row_removed(n, false);
         std::vector<bool> col_removed(m, false);
-        
+
         struct Assignment {
             int col;
             int row_idx;
         };
         std::vector<Assignment> assignment_stack;
         assignment_stack.reserve(n);
-        
+
         int head = 0;
         while(head < peel_stack.size()) {
             int j = peel_stack[head++];
@@ -173,15 +173,15 @@ extern "C" {
         for(int i=(int)assignment_stack.size()-1; i>=0; --i) {
             auto a = assignment_stack[i];
             const auto& e = edges[a.row_idx];
-            
+
             __m128i val1 = _mm_loadu_si128(&P_local[e.h1]);
             __m128i val2 = _mm_loadu_si128(&P_local[e.h2]);
             __m128i val3 = _mm_loadu_si128(&P_local[e.h3]);
             __m128i target = vals[e.key_idx];
-            
+
             __m128i current = _mm_xor_si128(_mm_xor_si128(val1, val2), val3);
             __m128i diff = _mm_xor_si128(target, current);
-            
+
             _mm_storeu_si128(&P_local[a.col], diff);
         }
         return true;
@@ -189,15 +189,15 @@ extern "C" {
 
     void solve_okvs_opt(uint64_t* keys, uint64_t* values, uint64_t* output, uint64_t n, uint64_t m, uint64_t* seed_ptr) {
         __m128i seed = _mm_loadu_si128((__m128i*)seed_ptr);
-        
+
         // 1. Calculate Bin Boundaries
         // We divide M evenly among bins. The remainder is distributed to the first few bins.
         std::vector<uint64_t> bin_offsets(NUM_BINS + 1);
         std::vector<uint64_t> m_per_bin(NUM_BINS);
-        
+
         uint64_t base_m = m / NUM_BINS;
         uint64_t remainder = m % NUM_BINS;
-        
+
         uint64_t current_offset = 0;
         for(uint64_t b=0; b<NUM_BINS; ++b) {
             bin_offsets[b] = current_offset;
@@ -208,18 +208,18 @@ extern "C" {
 
         // 2. Partition Data (Stateless)
         // Note on "Two-Choice Hashing":
-        // While Two-Choice Hashing (selecting the lighter of 2 potential bins) would significantly 
+        // While Two-Choice Hashing (selecting the lighter of 2 potential bins) would significantly
         // reduce max bin load variance, it introduces "Statefulness".
         // The bin assignment for Key K would depend on the load of bins, which depends on other keys.
-        // In standard PSI protocols (like RR22), the Decode step must be capable of processing keys 
+        // In standard PSI protocols (like RR22), the Decode step must be capable of processing keys
         // independently or without knowledge of the full set distribution (Sender/Receiver separation).
         // Therefore, we use **Simple Binning** (Stateless Hash) where Bin(K) = H(K) % Bins.
-        // We mitigate the resulting variance ("Balls-in-Bins" problem) by using a slightly larger 
+        // We mitigate the resulting variance ("Balls-in-Bins" problem) by using a slightly larger
         // expansion factor (epsilon ~ 1.35) which is bandwidth-acceptable and ensures stability.
-        
+
         std::vector<std::vector<uint64_t>> bin_keys(NUM_BINS);
         std::vector<std::vector<__m128i>> bin_vals(NUM_BINS);
-        
+
         // Pre-allocate to reduce reallocation overhead (assume ~uniform distribution)
         // 1.5x margin for pre-allocation safety
         size_t est_size = (n / NUM_BINS) * 3 / 2;
@@ -227,14 +227,14 @@ extern "C" {
             bin_keys[b].reserve(est_size);
             bin_vals[b].reserve(est_size);
         }
-        
+
         const __m128i* V_ptr = (const __m128i*)values;
         for(uint64_t i=0; i<n; ++i) {
             uint64_t b = get_bin_index(keys[i], seed);
             bin_keys[b].push_back(keys[i]);
             bin_vals[b].push_back(_mm_loadu_si128(&V_ptr[i]));
         }
-        
+
         // 3. Parallel Solve
         // Each bin is solved independently. This logic is perfectly parallelizable (embarrassingly parallel).
         // The working set for each bin (~1000 items) stays hot in L1 Cache.
@@ -244,10 +244,10 @@ extern "C" {
         #pragma omp parallel for schedule(dynamic)
         for(uint64_t b=0; b<NUM_BINS; ++b) {
             if(bin_keys[b].empty()) continue;
-            
+
             uint64_t offset = bin_offsets[b];
             uint64_t valid_m = m_per_bin[b];
-            
+
             if(!solve_bin(bin_keys[b], bin_vals[b], &P_vec[offset], valid_m, seed)) {
                 // On failure, log and fall back to the robust solver for this bin.
                 // The fallback is executed inside a critical section to avoid nested OpenMP
@@ -259,7 +259,7 @@ extern "C" {
 
                     // Prepare pointers for the safe solver
                     uint64_t* keys_ptr = &bin_keys[b][0];
-                    uint64_t* vals_ptr = &bin_vals[b][0];
+                    uint64_t* vals_ptr = (uint64_t*)&bin_vals[b][0];  // Cast __m128i* to uint64_t*
                     uint64_t* out_ptr = output + (offset * 2ULL); // each 128-bit slot == 2 uint64_t
 
                     // Call the safe solver implemented in okvs.cpp
@@ -273,10 +273,10 @@ extern "C" {
         __m128i seed = _mm_loadu_si128((__m128i*)seed_ptr);
         __m128i* P_vec = (__m128i*)storage;
         __m128i* out_vec = (__m128i*)output;
-        
+
         // Replicate Boundary Logic
         std::vector<uint64_t> bin_offsets(NUM_BINS + 1);
-        std::vector<uint64_t> m_per_bin(NUM_BINS); 
+        std::vector<uint64_t> m_per_bin(NUM_BINS);
         uint64_t base_m = m / NUM_BINS;
         uint64_t remainder = m % NUM_BINS;
         uint64_t current_offset = 0;
@@ -290,12 +290,12 @@ extern "C" {
         #pragma omp parallel for schedule(static)
         for(uint64_t i=0; i<n; ++i) {
             uint64_t b = get_bin_index(keys[i], seed);
-             
+
             uint64_t m_local = m_per_bin[b];
             uint64_t offset = bin_offsets[b];
-             
+
             Indices idx = get_bin_local_indices(keys[i], m_local, seed);
-             
+
             __m128i val = _mm_xor_si128(
                 _mm_xor_si128(_mm_loadu_si128(&P_vec[offset + idx.h1]), _mm_loadu_si128(&P_vec[offset + idx.h2])),
                 _mm_loadu_si128(&P_vec[offset + idx.h3])

mplang/{v2/kernels → kernels}/py_kernels.py

@@ -23,7 +23,7 @@ from __future__ import annotations
 
 import numpy as np
 
-from mplang.v2.libs.mpc.common.constants import (
+from mplang.libs.mpc.common.constants import (
     GOLDEN_RATIO_64,
     SPLITMIX64_GAMMA_2,
     SPLITMIX64_GAMMA_3,

mplang/{v2/libs → libs}/collective.py

@@ -30,8 +30,8 @@ Naming Convention:
 - collect: gather (N parties → 1 party, stacked)
 
 Example:
-    >>> from mplang.v2.libs.collective import transfer, replicate, distribute, collect
-    >>> from mplang.v2.dialects.simp import constant, converge
+    >>> from mplang.libs.collective import transfer, replicate, distribute, collect
+    >>> from mplang.dialects.simp import constant, converge
     >>>
     >>> # Create data on party 0
     >>> x = constant((0,), 42)
@@ -47,9 +47,9 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from mplang.v2.dialects.simp import converge, shuffle_static
-from mplang.v2.edsl import Object
-from mplang.v2.edsl.typing import MPType
+from mplang.dialects.simp import converge, shuffle_static
+from mplang.edsl import Object
+from mplang.edsl.typing import MPType
 
 if TYPE_CHECKING:
     pass

mplang/{v2/libs → libs}/device/__init__.py

@@ -17,7 +17,7 @@
 This module provides the high-level device-centric programming interface.
 """
 
-from mplang.v2.dialects.tensor import jax_fn
+from mplang.dialects.tensor import jax_fn
 
 from .api import (
     DeviceContext,

mplang/{v2/libs → libs}/device/api.py

@@ -28,10 +28,10 @@ from typing import Any, cast
 
 from jax.tree_util import tree_flatten, tree_map
 
-from mplang.v2.backends import load_builtins
-from mplang.v2.dialects import crypto, simp, spu, tee
-from mplang.v2.edsl.object import Object
-from mplang.v2.libs.device.cluster import Device
+from mplang.backends import load_builtins
+from mplang.dialects import crypto, simp, spu, tee
+from mplang.edsl.object import Object
+from mplang.libs.device.cluster import Device
 
 load_builtins()
 
@@ -43,7 +43,7 @@ def _resolve_cluster() -> Any:
     Interpreter with a _cluster_spec attribute. This allows nested contexts
     to override the cluster if needed.
     """
-    from mplang.v2.edsl.context import find_context
+    from mplang.edsl.context import find_context
 
     ctx = find_context(lambda c: getattr(c, "_cluster_spec", None) is not None)
     if ctx is not None:
@@ -356,7 +356,7 @@ class DeviceContext:
             if self._is_spu_device():
                 return self(fn)
             # PPU/TEE need tensor.jax_fn to compile JAX code
-            from mplang.v2.dialects.tensor import jax_fn
+            from mplang.dialects.tensor import jax_fn
 
             return self(jax_fn(fn))
 
@@ -443,7 +443,7 @@ def _ensure_tee_session(
     Returns:
         Tuple of (sess_frm, sess_tee) where each is a symmetric key Object
     """
-    import mplang.v2.edsl as el
+    import mplang.edsl as el
 
     # Get current context ID for cache isolation
     current_ctx = el.get_current_context()
@@ -749,11 +749,11 @@ def fetch(obj: Object) -> Any:
     Returns:
         Python value (numpy array, scalar, etc.)
     """
-    from mplang.v2.backends.simp_driver.state import SimpDriver
-    from mplang.v2.backends.simp_driver.values import DriverVar
-    from mplang.v2.edsl.context import get_current_context
-    from mplang.v2.runtime.interpreter import InterpObject, Interpreter
-    from mplang.v2.runtime.value import WrapValue
+    from mplang.backends.simp_driver.state import SimpDriver
+    from mplang.backends.simp_driver.values import DriverVar
+    from mplang.edsl.context import get_current_context
+    from mplang.runtime.interpreter import InterpObject, Interpreter
+    from mplang.runtime.value import WrapValue
 
     def _unwrap_value(val: Any) -> Any:
         """Unwrap WrapValue to get the underlying data."""