JAX数组动态切片索引报错及静态参数尝试失败问题

问题背景 #

需要为darkField数组每个元素生成2D高斯补丁，补丁尺寸理想计算公式为2 * np.ceil(3 * sigma) + 1（sigma为darkField对应元素值），示例中固定为10以规避错误。将高斯补丁归一化后，乘以intensityRefracted2DF对应元素得到模糊效果，最终添加到intensityRefracted3数组中。但运行代码时出现索引错误，尝试将索引设为静态参数后又出现哈希错误。

报错原因分析 #

1. IndexError：JIT编译函数中，无法使用动态起始/结束位置的普通切片索引（如i - size2:i + size2 + 1这类依赖动态值的切片），必须使用JAX提供的jax.lax.dynamic_update_slice处理动态位置的切片更新。
2. ValueError：vmap处理后的索引是BatchTracer类型的动态批量数据，无法作为静态参数传递给JIT函数——静态参数要求是可哈希的静态值，动态批量数据不符合该要求。

解决方案 #

• 用jax.lax.dynamic_update_slice替代普通切片的at.add操作，实现动态位置的数组更新。
• 移除apply_dark_field的static_argnums设置，因为vmap传递的是动态索引，不适合作为静态参数。
• 对vmap返回的批量数组进行求和，合并所有补丁的贡献（原代码中vmap返回的是每个索引对应的独立数组，需累加得到最终结果）。

修改后的完整代码 #

import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
from functools import partial

@partial(jax.jit, static_argnums=(1,))
def gaussian_shape(sigma, size):
    """
    Generate a Gaussian shape.

    Args:
        sigma (float or 2D numpy array): Standard deviation(s) of the Gaussian shape.
        size (int): Size of the Gaussian shape.

    Returns:
        exponent (2D numpy array): Gaussian shape.
    """
    x = jnp.arange(0, size) - jnp.floor(size / 2)
    exponent = jnp.exp(-(x ** 2) / (2 * sigma ** 2))
    exponent = jnp.outer(exponent, exponent)
    exponent /= jnp.sum(exponent)
    return exponent

@jax.jit
def apply_dark_field(i, j, intensityRefracted2DF, darkField):
    currDF_ij = darkField[i, j]
    patch_size = 10  # 后续可替换为动态计算：2 * jnp.ceil(3 * currDF_ij).astype(int) + 1
    patch = gaussian_shape(currDF_ij, patch_size)
    patch = patch * intensityRefracted2DF[i, j]
    
    # 创建初始全零数组，存放当前补丁的贡献
    curr_contribution = jnp.zeros_like(intensityRefracted2DF)
    # 计算切片起始位置
    start_idx = (i - patch_size//2, j - patch_size//2)
    # 用dynamic_update_slice将补丁放到对应位置
    curr_contribution = jax.lax.dynamic_update_slice(curr_contribution, patch, start_idx)
    return curr_contribution

@jax.jit
def darkFieldLoop(intensityRefracted2DF, darkField):
    # 仅处理intensityRefracted2DF非零的位置
    mask = intensityRefracted2DF != 0
    indices_i, indices_j = jnp.where(mask)
    
    # 对每个索引应用apply_dark_field，得到所有补丁的贡献数组
    contributions = jax.vmap(apply_dark_field, in_axes=(0, 0, None, None))(indices_i, indices_j, intensityRefracted2DF, darkField)
    
    # 累加所有贡献得到最终结果
    intensityRefracted3 = jnp.sum(contributions, axis=0)
    return intensityRefracted3

# 测试代码
intensityRefracted2DF = np.random.rand(10,10)
darkField = np.random.rand(10, 10)

a = darkFieldLoop(intensityRefracted2DF, darkField)

plt.imshow(a)
plt.show()

额外说明 #

如果需要实现动态补丁尺寸（根据每个sigma计算不同的补丁大小），需注意JAX中动态形状的处理，可结合jax.vmap与jax.lax.cond，或使用jax.jit的dynamic=True参数（需对应JAX版本支持）。但动态形状会带来一定性能损耗，静态尺寸在JIT编译下性能更优。

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