在JAX中高效实现可变长度参数化矩阵链与向量的乘法运算

核心思路 #

抛弃冗余的填充操作，利用JAX原生的序列处理能力，针对每个变长索引序列单独执行矩阵链乘，再通过批量映射完成多序列计算，同时保留自动微分的兼容性。

具体实现方案 #

方案1：单序列scan + 批量vmap（无填充，效率最优） #

先将字符串形式的门控序列转换为索引数组，对单个序列用scan完成矩阵乘，再通过vmap批量处理所有序列，完全避免无效计算。

修改后的完整代码：

import jax
import jax.numpy as jnp
from jax.scipy.linalg import expm

sigma0 = jnp.eye(2)
sigmax = jnp.array([[0, 1], [1, 0]])
sigmay = jnp.array([[0, -1j], [1j, 0]])
sigmaz = jnp.array([[1, 0], [0, -1]])
sigma = jnp.array([sigmax, sigmay, sigmaz])

def gates_func(params):
    theta = params["theta"]
    epsilon = params["epsilon"]

    n = jnp.array([jnp.cos(theta), 0, jnp.sin(theta)])
    omega = jnp.pi / 2 * (1 + epsilon)
    X90 = expm(-1j * omega * jnp.einsum("i,ijk->jk", n, sigma) / 2)

    # 返回按索引映射的门控数组：0对应X90，1对应Z90
    return jnp.array([X90, expm(-1j * jnp.pi / 2 * sigmaz / 2)])

# 处理单个序列的函数：输入门控数组、序列索引、初始向量，返回最终结果
def process_single_seq(gates, seq_indices, init_vec):
    def step(carry, idx):
        return gates[idx] @ carry, None
    final_vec, _ = jax.lax.scan(step, init_vec, seq_indices)
    return jnp.abs(final_vec[0])

def multiply_out(params):
    gates = gates_func(params)
    # 将字符串序列转为索引数组
    gate_indices_list = [
        jnp.array([0, 0]),
        jnp.array([0, 1]),
        jnp.array([1, 0]),
        jnp.array([0, 1, 0])
    ]
    init_vec = jnp.array([1.0, 0.0], dtype=jnp.complex128)
    
    # 批量处理所有序列
    results = jnp.array([process_single_seq(gates, seq, init_vec) for seq in gate_indices_list])
    return results

# 测试计算与梯度
params = dict(theta=-0.0, epsilon=0.001)
print(multiply_out(params))
grad_fn = jax.grad(multiply_out)
print(grad_fn(params))

方案2：掩码式填充（兼容统一批量处理） #

如果需要统一的张量格式，可通过掩码标记有效计算步骤，跳过填充位置的无效操作：

def multiply_out_masked(params):
    gates = gates_func(params)
    gate_indices_list = [
        jnp.array([0, 0]),
        jnp.array([0, 1]),
        jnp.array([1, 0]),
        jnp.array([0, 1, 0])
    ]
    max_len = max(len(seq) for seq in gate_indices_list)
    # 填充序列并生成掩码（1为有效步骤，0为填充）
    padded_seqs = jnp.array([jnp.pad(seq, (0, max_len - len(seq)), mode="constant", constant_values=-1) for seq in gate_indices_list])
    masks = jnp.array([jnp.pad(jnp.ones(len(seq)), (0, max_len - len(seq)), mode="constant") for seq in gate_indices_list])
    init_vec = jnp.array([1.0, 0.0], dtype=jnp.complex128)
    
    def step(carry, inputs):
        seq_idx, mask = inputs
        # 掩码为真时执行矩阵乘，否则保持当前值
        new_carry = jax.lax.select(mask.astype(bool), gates[seq_idx] @ carry, carry)
        return new_carry, None
    
    # 批量扫描所有序列
    final_vecs, _ = jax.vmap(lambda seq, mask: jax.lax.scan(step, init_vec, (seq, mask))[0])(padded_seqs, masks)
    results = jnp.abs(final_vecs[:, 0])
    return results

方案对比 #

• 方案1完全无填充，每个序列仅计算实际长度的步骤，适配长尾分布场景，计算效率最高
• 两种方案均保留JAX自动微分能力，可正常对参数p求梯度
• 代码结构简洁，易扩展门控类型或新增序列

内容的提问来源于stack exchange，提问作者Evan