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

mplang/v2/libs/mpc/ot/extension.py

@@ -0,0 +1,477 @@
+# 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.
+
+"""OT Extension (IKNP).
+
+Implements IKNP OT extension protocol to perform N OTs using k Base OTs.
+Ref: https://crypto.stanford.edu/~valeria/research/2003/IKNP03.pdf
+"""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import jax.numpy as jnp
+
+import mplang.v2.edsl as el
+from mplang.v2.dialects import crypto, field, simp, tensor
+
+
+def prg_expand(seed_tensor: el.Object, length: int) -> el.Object:
+    """Pseudo-Random Generator: Expand seed to `length` bits (as uint8 0/1).
+
+    Uses AES-NI via field.aes_expand for cryptographic security.
+    """
+    # Calculate number of 128-bit blocks needed to cover `length` bits.
+    # field.aes_expand returns (K, M, 2) uint64 blocks.
+    # Total bits = M * 128.
+
+    m_blocks = (length + 127) // 128
+
+    # Input seed_tensor is (K, 32) bytes (uint8).
+    # field.aes_expand expects (K, 2) uint64 seeds.
+
+    def _reshape_seeds(s_bytes: Any) -> Any:
+        # s_bytes: (K, 32) u8
+        # Take first 16 bytes for 128-bit key/seed
+        s_16 = s_bytes[:, :16]
+        return s_16.view(jnp.uint64).reshape(-1, 2)
+
+    seeds_u64 = tensor.run_jax(_reshape_seeds, seed_tensor)
+
+    expanded_blocks = field.aes_expand(seeds_u64, m_blocks)  # (K, M, 2) u64
+
+    # Convert blocks to bits
+    def _blocks_to_bits(blocks: Any) -> Any:
+        # blocks: (K, M, 2) u64
+        # unpackbits
+        # view as u8
+        bytes_view = blocks.view(jnp.uint8)  # (K, M, 16)
+        bits = jnp.unpackbits(bytes_view, axis=-1, bitorder="little")  # (K, M, 128)
+
+        # Flatten last two dims
+        bits_flat = bits.reshape(bits.shape[0], -1)
+
+        # Slice to exact length
+        return bits_flat[:, :length]
+
+    return cast(el.Object, tensor.run_jax(_blocks_to_bits, expanded_blocks))
+
+
+def vec_hash(data_bytes: el.Object, domain_sep: int, num_rows: int) -> el.Object:
+    """Hash rows of a (N, D) tensor independently.
+
+    Args:
+        data_bytes: (N, D) tensor to hash.
+        domain_sep: Integer domain separator to mix into the hash.
+        num_rows: Number of rows N. Must be provided explicitly.
+    """
+    # Optimized batch hashing:
+    # 1. Prepend domain_sep to all rows (vectorized concatenation)
+    # 2. Call crypto.hash_bytes once on the whole tensor
+
+    if domain_sep != 0:
+
+        def _prepend_ds(arr: Any, ds: int) -> Any:
+            # arr: (N, D)
+            N = arr.shape[0]
+            # Create (N, 8) domain sep block using repeat & reshape
+            # ds_arr: (8,)
+            ds_arr = jnp.array([ds], dtype=jnp.uint64).view(jnp.uint8)
+            # Broadcast to (N, 8)
+            ds_block = jnp.broadcast_to(ds_arr, (N, 8))
+
+            return jnp.concatenate([ds_block, arr], axis=1)
+
+        # Result: (N, D+8)
+        data_to_hash = tensor.run_jax(lambda a: _prepend_ds(a, domain_sep), data_bytes)
+    else:
+        data_to_hash = data_bytes
+
+    # Call batched hash_bytes
+    # Input: (N, D_total) -> Output: (N, 32)
+    # This generates a single graph node, solving the compiler explosion issue.
+    # explicit hash_batch primitive (rank >= 2)
+    hashes = crypto.hash_batch(data_to_hash)
+
+    return hashes
+
+
+def iknp_core(
+    choice_bits: el.Object, sender: int, receiver: int, num_ots: int
+) -> tuple[el.Object, el.Object, el.Object]:
+    """Core IKNP Matrix Generation.
+
+    Returns:
+        t_matrix: (N, K) bit matrix on Sender.
+        q_matrix: (N, K) bit matrix on Receiver.
+        s_choices: (K,) choice bits on Sender (s).
+    """
+    K = 128
+
+    # 1. Base OTs
+    def gen_s() -> el.Object:
+        # Generate random bits at runtime using new API
+        return crypto.random_bits(K)
+
+    s = simp.pcall_static((sender,), gen_s)
+
+    def gen_seeds() -> tuple[el.Object, el.Object]:
+        # Generate random bytes at runtime
+        k0_bytes = crypto.random_bytes(K * 32)
+        k1_bytes = crypto.random_bytes(K * 32)
+
+        # Reshape to (K, 32) using run_jax for XLA optimization
+        def _reshape_k32(b: Any) -> Any:
+            return b.reshape(K, 32)
+
+        k0 = tensor.run_jax(_reshape_k32, k0_bytes)
+        k1 = tensor.run_jax(_reshape_k32, k1_bytes)
+        return k0, k1
+
+    k0_base, k1_base = simp.pcall_static((receiver,), gen_seeds)
+
+    # Base OT Logic (Inlined)
+    # C (Common Point) initialization
+    # SECURITY FIX: C must be generated by the Base Sender (receiver in IKNP context)
+    # to prevent Base Receiver (sender in IKNP context) from knowing the discrete log,
+    # which would allow them to decrypt both messages and recover choice bits.
+    def base_param_gen() -> el.Object:
+        C = crypto.ec_mul(crypto.ec_generator(), crypto.ec_random_scalar())
+        return C
+
+    C_point = simp.pcall_static((receiver,), base_param_gen)
+    C_for_sender = simp.shuffle_static(C_point, {sender: receiver})
+
+    # Duplicate initialization removed
+
+    # R (Sender of BaseOT) keygen
+    def base_sender_keygen(
+        C: el.Object, s_base_choices: el.Object
+    ) -> tuple[el.Object, list[el.Object]]:
+        # s_base_choices is (K,) Tensor
+        PK0_bytes_list = []
+        k_priv_list = []
+
+        for i in range(K):
+            k_priv = crypto.ec_random_scalar()
+            PK_sigma = crypto.ec_mul(crypto.ec_generator(), k_priv)
+
+            # Slice s[i]
+            # Using slice_tensor here is efficient since s_base_choices is small; the overhead of run_jax is unnecessary.
+            s_i = tensor.slice_tensor(s_base_choices, (i,), (i + 1,))
+            s_scalar = crypto.ec_scalar_from_int(s_i)
+
+            diff = crypto.ec_sub(C, PK_sigma)
+            # select checks s_scalar. If 1 (true), pick diff.
+            PK0 = crypto.select(s_scalar, diff, PK_sigma)
+
+            # Convert to bytes (65 bytes uncompressed)
+            # K is 128, so overhead is small.
+            # Stacking points for shuffle.
+            pk0_b = crypto.ec_point_to_bytes(PK0)
+            # Reshape for stack: (65,) -> (1, 65)
+            pk0_b_r = tensor.reshape(pk0_b, (1, 65))
+
+            PK0_bytes_list.append(pk0_b_r)
+            k_priv_list.append(k_priv)
+
+        # Stack into (K, 65)
+        PK0_stacked = tensor.concat(PK0_bytes_list, axis=0)
+
+        return PK0_stacked, k_priv_list
+
+    # base_keys -> (PK0_stacked, k_priv_list (TraceObject list))
+    # Pass C_for_sender (received from receiver) to sender
+    base_keys_tuple = simp.pcall_static((sender,), base_sender_keygen, C_for_sender, s)
+
+    # Extract
+    PK0_loc = simp.pcall_static((sender,), lambda x: x[0], base_keys_tuple)
+    # Note: k_priv (x[1]) stays on sender, used later in base_decrypt_rev via base_keys_tuple
+    PK0_recv = simp.shuffle_static(PK0_loc, {receiver: sender})
+
+    # R (Base Sender) Encrypts k0, k1
+    def base_encrypt_rev(
+        C: el.Object,
+        PK0_bytes_tensor: el.Object,
+        m0_tensor: el.Object,
+        m1_tensor: el.Object,
+    ) -> tuple[el.Object, el.Object, el.Object]:
+        # m0, m1 are (K, 32) tensors.
+        # PK0_bytes_tensor is (K, 65)
+
+        U_bytes_list = []
+        c0_list = []
+        c1_list = []
+
+        for i in range(K):
+            # Unstack PK0
+            # PK0_bytes_tensor is (K, 65)
+            # We want row i, all 65 bytes: slice(i:i+1, 0:65)
+            pk0_b = tensor.slice_tensor(PK0_bytes_tensor, (i, 0), (i + 1, 65))
+            # Reshape to (65,) for conversion
+            pk0_b_flat = tensor.reshape(pk0_b, (65,))
+            PK0 = crypto.ec_bytes_to_point(pk0_b_flat)
+
+            r = crypto.ec_random_scalar()
+            U = crypto.ec_mul(crypto.ec_generator(), r)
+
+            # Stack U as bytes
+            u_b = crypto.ec_point_to_bytes(U)
+            u_b_r = tensor.reshape(u_b, (1, 65))
+            U_bytes_list.append(u_b_r)
+
+            K0_point = crypto.ec_mul(PK0, r)
+            PK1 = crypto.ec_sub(C, PK0)
+            K1_point = crypto.ec_mul(PK1, r)
+
+            sk0 = crypto.hash_bytes(crypto.ec_point_to_bytes(K0_point))  # (32,)
+            sk1 = crypto.hash_bytes(crypto.ec_point_to_bytes(K1_point))
+
+            # Extract row i and encrypt in single run_jax block
+            def _slice_and_enc(
+                m0_full: Any, m1_full: Any, k0: Any, k1: Any, idx: int = i
+            ) -> tuple[Any, Any]:
+                # Slice row i and reshape to (32,)
+                m0_row = m0_full[idx].flatten()
+                m1_row = m1_full[idx].flatten()
+                # XOR with keys
+                c0 = jnp.bitwise_xor(m0_row, k0)
+                c1 = jnp.bitwise_xor(m1_row, k1)
+                return c0.reshape(1, -1), c1.reshape(1, -1)
+
+            c0, c1 = tensor.run_jax(_slice_and_enc, m0_tensor, m1_tensor, sk0, sk1)
+
+            c0_list.append(c0)
+            c1_list.append(c1)
+
+        # Stack outputs
+        U_stacked = tensor.concat(U_bytes_list, axis=0)  # (K, 65)
+        c0_stacked = tensor.concat(c0_list, axis=0)  # (K, 32) (assuming 32 byte msgs)
+        c1_stacked = tensor.concat(c1_list, axis=0)  # (K, 32)
+
+        return U_stacked, c0_stacked, c1_stacked
+
+    base_cts_rev = simp.pcall_static(
+        (receiver,), base_encrypt_rev, C_point, PK0_recv, k0_base, k1_base
+    )
+
+    # Shuffle tuple(Tensor, Tensor, Tensor) - Efficient!
+    # tree_map handles tuple
+    from jax.tree_util import tree_map
+
+    base_cts_s = tree_map(
+        lambda x: simp.shuffle_static(x, {sender: receiver}), base_cts_rev
+    )
+
+    def base_decrypt_rev(
+        keys: tuple[el.Object, list[el.Object]],
+        cts: tuple[el.Object, el.Object, el.Object],
+        s_choices: el.Object,
+    ) -> el.Object:
+        # keys[0] is PK0_stacked (unused here), keys[1] is k_priv_list
+        _, k_priv_list = keys
+        # cts are stacked (K, 65), (K, 32), (K, 32)
+        U_packed, c0_packed, c1_packed = cts
+
+        k_s_rows = []
+
+        for i in range(K):
+            k_priv = k_priv_list[i]
+
+            # Unstack U
+            u_b = tensor.slice_tensor(U_packed, (i, 0), (i + 1, 65))
+            u_b_flat = tensor.reshape(u_b, (65,))
+            U = crypto.ec_bytes_to_point(u_b_flat)
+
+            # Unstack c0, c1
+            c0 = tensor.slice_tensor(c0_packed, (i, 0), (i + 1, 32))
+            c1 = tensor.slice_tensor(c1_packed, (i, 0), (i + 1, 32))
+            # Reshape to (32,) or (1, 32) depending on what _slice_dec_reshape expects
+            # _slice_dec_reshape expects (32,) usually if we want flat XOR?
+            # Let's check _slice_dec_reshape: it does `jnp.bitwise_xor(chosen_c, k)`.
+            # sk is (32,). So chosen_c should be (32,).
+            c0 = tensor.reshape(c0, (32,))
+            c1 = tensor.reshape(c1, (32,))
+
+            # Recov K = U^k_priv
+            SharedK = crypto.ec_mul(U, k_priv)
+            sk = crypto.hash_bytes(crypto.ec_point_to_bytes(SharedK))
+
+            # Combined slice + decrypt + reshape in single run_jax
+            def _slice_dec_reshape(
+                s_arr: Any, k: Any, c0_: Any, c1_: Any, idx: int = i
+            ) -> Any:
+                sel = s_arr[idx]
+                chosen_c = jnp.where(sel == 0, c0_, c1_)
+                result = jnp.bitwise_xor(chosen_c, k)
+                return result.reshape(1, 32)
+
+            res_row = tensor.run_jax(_slice_dec_reshape, s_choices, sk, c0, c1)
+            k_s_rows.append(res_row)
+
+        # Concat using tensor.concat (run_jax with many args can cause tracing issues)
+        return tensor.concat(k_s_rows, axis=0)
+
+    k_s = simp.pcall_static((sender,), base_decrypt_rev, base_keys_tuple, base_cts_s, s)
+
+    # 2. PRG Expansion & Correction
+    def calc_u(k0_loc: el.Object, k1_loc: el.Object, r_loc: el.Object) -> el.Object:
+        g_k0 = prg_expand(k0_loc, num_ots)  # (K, num_ots)
+        g_k1 = prg_expand(k1_loc, num_ots)  # (K, num_ots)
+
+        # choice_bits can be:
+        # - (N,) 1D vector for standard IKNP
+        # - (N, K) 2D matrix for KKRT OPRF
+        #
+        # For IKNP: u^j = G(k0^j) ^ G(k1^j) ^ r, where r is broadcast to all K rows
+        # For KKRT: u^j = G(k0^j) ^ G(k1^j) ^ r^j, where r is (N, K) transposed to (K, N)
+
+        # Handle both 1D and 2D inputs
+        def _compute_u(g0: Any, g1: Any, r: Any) -> Any:
+            # g0, g1: (K, N) bit matrices
+            # r: either (N,) or (N, K)
+            if r.ndim == 1:
+                # 1D case: broadcast (N,) -> (1, N) for XOR with (K, N)
+                r_t = jnp.expand_dims(r, axis=0)  # (1, N)
+            else:
+                # 2D case: transpose (N, K) -> (K, N)
+                r_t = jnp.transpose(r, (1, 0))  # (K, N)
+            return jnp.bitwise_xor(jnp.bitwise_xor(g0, g1), r_t)
+
+        return cast(el.Object, tensor.run_jax(_compute_u, g_k0, g_k1, r_loc))
+
+    u = simp.pcall_static((receiver,), calc_u, k0_base, k1_base, choice_bits)
+    u_recv = simp.shuffle_static(u, {sender: receiver})
+
+    # 3. Matrix Recovery & Transpose
+    def calc_t(k_s_loc: el.Object, u_loc: el.Object, s_loc: el.Object) -> el.Object:
+        g_k_s = prg_expand(k_s_loc, num_ots)
+
+        def _recover_and_transpose(g: Any, mask: Any, sel: Any) -> Any:
+            # Combine recover and transpose into single XLA block
+            sel_exp = jnp.expand_dims(sel, axis=-1)
+            term = jnp.bitwise_and(mask, sel_exp)
+            t_rows = jnp.bitwise_xor(g, term)
+            return jnp.transpose(t_rows, (1, 0))  # (N, K)
+
+        return cast(
+            el.Object, tensor.run_jax(_recover_and_transpose, g_k_s, u_loc, s_loc)
+        )
+
+    t_matrix = simp.pcall_static((sender,), calc_t, k_s, u_recv, s)
+
+    def calc_q(k0_loc: el.Object) -> el.Object:
+        g_k0 = prg_expand(k0_loc, num_ots)
+        # Use run_jax for transpose to enable XLA fusion
+        return cast(el.Object, tensor.run_jax(lambda x: jnp.transpose(x, (1, 0)), g_k0))
+
+    q_matrix = simp.pcall_static((receiver,), calc_q, k0_base)
+
+    # s is on Sender. t_matrix on Sender. q_matrix on Receiver.
+    return t_matrix, q_matrix, s
+
+
+def s_choices_sender(s: el.Object) -> el.Object:
+    return s  # Already pcalled on sender
+
+
+def transfer_extension(
+    m0: el.Object,
+    m1: el.Object,
+    choice_bits: el.Object,
+    sender: int,
+    receiver: int,
+    num_ots: int,
+) -> el.Object:
+    """Perform IKNP OT Extension."""
+
+    t_matrix, q_matrix, s = iknp_core(choice_bits, sender, receiver, num_ots)
+
+    # 4. Encryption
+    def encrypt_msgs(
+        t_loc: el.Object, s_loc: el.Object, m0_loc: el.Object, m1_loc: el.Object
+    ) -> el.Object:
+        # t: (N, K)
+        # s: (K,)
+
+        # Hash keys before using them as masks to break linear correlation
+        # H(t) and H(t^s)
+        # We use domain_sep=1 for IKNP payload masking
+        h_t = vec_hash(t_loc, domain_sep=1, num_rows=num_ots)
+
+        def _xor_s_and_hash(t: Any, s: Any) -> Any:
+            t_xor_s = jnp.bitwise_xor(t, s)
+            return t_xor_s
+
+        # We need to compute H(t^s). We can't easily do it in one block with vec_hash
+        # unless we compute t^s first.
+        t_xor_s_loc = cast(el.Object, tensor.run_jax(_xor_s_and_hash, t_loc, s_loc))
+        h_t_xor_s = vec_hash(t_xor_s_loc, domain_sep=1, num_rows=num_ots)
+
+        def _enc(ht: Any, hts: Any, msg0: Any, msg1: Any) -> Any:
+            # ht, hts are mapped to (N, 32) bytes usually, or whatever vec_hash returns
+            # msg0, msg1 are (N, D) bytes
+
+            # Ensure shapes match for XOR
+            # vec_hash returns (N, 32)
+            # If messages are not 32 bytes, we might need to adjust or truncation?
+            # Standard IKNP assumes messages are block size (128 bit = 16 bytes).
+            # But vec_hash produces 32 bytes (SHA256 usually).
+            # We slice hash to message length.
+
+            # msg0 shape: (N, 16) usually
+            d = msg0.shape[1]
+
+            ht_sliced = ht[:, :d]
+            hts_sliced = hts[:, :d]
+
+            c0 = jnp.bitwise_xor(msg0, ht_sliced)
+            c1 = jnp.bitwise_xor(msg1, hts_sliced)
+            return c0, c1
+
+        return cast(el.Object, tensor.run_jax(_enc, h_t, h_t_xor_s, m0_loc, m1_loc))
+
+    ciphertexts = simp.pcall_static((sender,), encrypt_msgs, t_matrix, s, m0, m1)
+
+    from jax.tree_util import tree_map
+
+    ciphertexts_recv = tree_map(
+        lambda x: simp.shuffle_static(x, {receiver: sender}), ciphertexts
+    )
+
+    def decrypt_msg(
+        q_loc: el.Object, r_loc: el.Object, c_texts: tuple[el.Object, el.Object]
+    ) -> el.Object:
+        c0, c1 = c_texts
+
+        # Hash q: H(q)
+        h_q = vec_hash(q_loc, domain_sep=1, num_rows=num_ots)
+
+        def _dec(hq: Any, r: Any, ct0: Any, ct1: Any) -> Any:
+            # hq: (N, 32)
+            d = ct0.shape[1]
+            hq_sliced = hq[:, :d]
+
+            m0_cand = jnp.bitwise_xor(ct0, hq_sliced)
+            m1_cand = jnp.bitwise_xor(ct1, hq_sliced)
+            r_exp = jnp.expand_dims(r, axis=-1)
+            return jnp.where(r_exp == 1, m1_cand, m0_cand)
+
+        return cast(el.Object, tensor.run_jax(_dec, h_q, r_loc, c0, c1))
+
+    res = simp.pcall_static(
+        (receiver,), decrypt_msg, q_matrix, choice_bits, ciphertexts_recv
+    )
+    return cast(el.Object, res)

mplang/v2/libs/mpc/ot/silent.py

@@ -0,0 +1,217 @@
+# 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.
+
+"""Silent OT (Random VOLE) Implementation.
+
+Implements "Silent Random VOLE" via Linear Expansion (LPN-like).
+This provides O(N) local computation but O(k) communication.
+"""
+
+from typing import Any, cast
+
+import jax.numpy as jnp
+
+import mplang.v2.edsl as el
+import mplang.v2.edsl.typing as elt
+import mplang.v2.libs.mpc.vole.gilboa as vole
+from mplang.v2.dialects import crypto, field, simp, tensor
+
+
+def silent_vole_random_u(
+    sender: int,
+    receiver: int,
+    n: int,
+    base_k: int = 1024,
+) -> tuple[el.Object, el.Object, el.Object, el.Object]:
+    """Execute Silent Random VOLE (Linear Expansion).
+
+    Args:
+        sender: Rank of Sender.
+        receiver: Rank of Receiver.
+        n: Target vector length (e.g. 10^9).
+        base_k: Size of Base VOLE (LPN parameter).
+
+    Returns:
+        v, w, u, delta
+        Where w = v + u * delta.
+        u is RANDOM.
+    """
+
+    # 1. Base VOLE (Standard Gilboa)
+    # We need providers for base_u and base_delta.
+
+    def _base_u_provider() -> el.Object:
+        # Random U_base (base_k, 2) using new API
+        return crypto.random_tensor((base_k, 2), elt.u64)
+
+    def _base_delta_provider() -> el.Object:
+        # Random Delta (2,) using new API
+        return crypto.random_tensor((2,), elt.u64)
+
+    # v_base: (k, 2), w_base: (k, 2)
+    # The return type is a Union, mypy complains about unpacking.
+    # We ignore the type error here as we know return_secrets=True returns 4 values.
+    v_base, w_base, u_base, delta = vole.vole(  # type: ignore
+        sender,
+        receiver,
+        base_k,
+        _base_u_provider,
+        _base_delta_provider,
+        return_secrets=True,
+    )
+
+    # 2. Linear Expansion
+    # We rely on a public seed for the mixing matrix M.
+    seed = simp.pcall_static((sender,), lambda: crypto.random_bytes(32))
+    # Share seed (Receiver needs it too)
+    seed_recv = simp.shuffle_static(seed, {receiver: sender})  # S -> R
+
+    # Expansion Logic
+    # We process in chunks to avoid massive implementation limit or memory issues.
+    # But EDSL graph optimization might handle loops?
+    # For safe side, let's implement a loop over chunks in Python if N is large.
+    # However, N is dynamic usually? Here N is int param.
+
+    # Chunk size
+    CHUNK_SIZE = 100_000  # 100k items per chunk
+
+    # We need to broadcast delta and bases to expansion function?
+    # Actually, we can just expand v_base -> v_long, w_base -> w_long.
+    # u_long is implicit (u_base * M).
+    # Since we need to return u, we expand u_base too.
+
+    # Define expansion op
+    def _expand_chunk(base_vec: Any, mask_packed: Any, chunk_len: int) -> Any:
+        # base_vec: (K, 2) u64
+        # mask_packed: (K, blocks, 2) u64 (AES output)
+        # chunk_len: number of bits to extract
+
+        # 1. Unpack bits from mask_packed
+        # mask_packed is (K, blocks, 2) u64.
+        # View as u8: (K, blocks, 16)
+        mask_u8 = mask_packed.view(jnp.uint8)
+
+        # Unpack bits: (K, blocks, 16, 8)
+        bits = jnp.unpackbits(mask_u8, bitorder="little")
+
+        # Flatten to (K, total_bits)
+        bits_flat = bits.reshape(base_k, -1)
+
+        # Slice to chunk_len
+        mask_mat = bits_flat[:, :chunk_len].astype(jnp.uint64)
+
+        # Broadcast base_vec (K, 2)
+        # We want: out[c] = XOR_sum_j (base[j] * mask[j, c])
+
+        base_shuffled = base_vec.reshape(base_k, 1, 2)
+        mask_expanded = mask_mat.reshape(base_k, chunk_len, 1)
+
+        # term[j, c] = base[j] * mask[j, c]
+        # mask is 0 or 1 (uint64). Multiplication works as selection.
+        terms = base_shuffled * mask_expanded  # (K, chunk, 2)
+
+        # XOR Reduce over K
+        # Use simple loop or scan.
+        # terms: (K, chunk, 2)
+
+        def _xor_scan(carry: Any, x: Any) -> tuple[Any, Any]:
+            new_carry = jnp.bitwise_xor(carry, x)
+            return new_carry, None
+
+        # init: (chunk, 2) zeros
+        init_val = jnp.zeros((chunk_len, 2), dtype=jnp.uint64)
+
+        from jax import lax
+
+        res, _ = lax.scan(_xor_scan, init_val, terms)
+
+        return res
+
+    # 3. Orchestration
+    # We iterate chunks on Host? Or use `scan`?
+    # Host loop is easier for Memory management (Streaming).
+    # Return a "Lazy Object" or List of Objects?
+    # The signature `silent_vole` usually returns full Tensor.
+    # User requirement: "Silent OT" to reduce communications.
+    # If we return a full (N,) tensor, we solved bandwidth but not RAM.
+    # But for Phase 2 task "Protocol Upgrade", bandwidth is key.
+    # Phase 2 task "Streaming" handles RAM.
+    # So returning full Tensor is "okay" for now, although it might OOM 1B.
+    # Let's implement blocked execution and stack? No, that OOMs.
+
+    # We will implement `silent_vole_random_u` to return a `BigTensor` handle?
+    # Or just `el.Object` (which might be huge).
+    # Since we are in EDSL, the `el.Object` represents the *computation*.
+    # If we return a graph that produces (10^9,) tensor, the Evaluator might crash trying to allocate it.
+
+    # Let's just implement loop and return concatenated for now, assume 10^7-10^8 test case.
+    # For 10^9, we rely on Streaming Refactor later.
+
+    num_chunks = (n + CHUNK_SIZE - 1) // CHUNK_SIZE
+
+    def _run_expansion(b: Any, seed_val: Any) -> el.Object:
+        # b: base (K, 2)
+        # seed_val: (32,) u8
+
+        # 1. Derive K seeds from master seed using combined run_jax block
+        def _view_slice_reshape(b: Any) -> Any:
+            # View as u64, slice first row, then reshape for AES expand
+            u64_view = b.view(jnp.uint64).reshape(-1, 2)
+            master_seed = u64_view[:1]  # (1, 2)
+            return master_seed
+
+        master_seed = tensor.run_jax(_view_slice_reshape, seed_val)
+
+        # Expand to K seeds: (1, K, 2)
+        row_seeds_packed = field.aes_expand(master_seed, base_k)
+        # Reshape using run_jax for XLA optimization
+        row_seeds = tensor.run_jax(lambda x: x.reshape(base_k, 2), row_seeds_packed)
+
+        # Iterate chunks
+        local_res = []
+        for i in range(num_chunks):
+            this_len = min(CHUNK_SIZE, n - i * CHUNK_SIZE)
+
+            # Generate mask for this chunk using AES
+            # Need ceil(this_len / 128) blocks
+            num_blocks = (this_len + 127) // 128
+            mask_packed = field.aes_expand(row_seeds, num_blocks)
+
+            # We must use `tensor.run_jax` so logic runs on device
+            def _core(base: Any, mask: Any, this_len: int = this_len) -> Any:
+                return _expand_chunk(base, mask, this_len)
+
+            chunk_res = tensor.run_jax(_core, b, mask_packed)
+            local_res.append(chunk_res)
+
+        # Use run_jax for concat to enable XLA fusion
+        if len(local_res) == 1:
+            return cast(el.Object, local_res[0])
+
+        def _concat_chunks(*chunks: Any) -> Any:
+            return jnp.concatenate(chunks, axis=0)
+
+        return cast(el.Object, tensor.run_jax(_concat_chunks, *local_res))
+
+    # Execute on Sender
+    v_long = simp.pcall_static((sender,), _run_expansion, v_base, seed)
+    # Execute on Receiver
+    w_long = simp.pcall_static((receiver,), _run_expansion, w_base, seed_recv)
+
+    # U expansion
+    u_long = simp.pcall_static((sender,), _run_expansion, u_base, seed)
+
+    # Delta is scalar, reusable
+
+    return v_long, w_long, u_long, delta