如何以JAX JIT可编译方式隔离数组的子集？

我有若干需分析的数据集，每个数据集的x值和y值分别存储在两个数组中，数组规模通常较大（约10万及以上数据点）。这些数据会通过多个JIT兼容函数对y值执行各类计算，但这类计算在大数据集下扩展性较差。因此我希望编写一个函数先隔离数据子集，再传入后续函数（通过分块处理全量数据集节省时间）。我已有一个可工作的子集隔离函数，但它不兼容JIT，导致我不得不取消主函数的JIT装饰（因为所有函数都由这个JIT编译的主函数调用）。

演示脚本如下：

import jax
import numpy as np
import matplotlib.pyplot as plt

# enable 64-bit floating point representation in JAX
jax.config.update("jax_enable_x64", True)

def isolate_subset(indices: list) -> tuple:
    # compute difference between successive x values
    delta = x[1] - x[0]
    
    # get indices of start and stop of subset
    i_start = jax.numpy.min(jax.numpy.where(jax.numpy.isclose(x, x[indices[0]], atol=delta/2))[0])
    i_stop = jax.numpy.max(jax.numpy.where(jax.numpy.isclose(x, x[indices[1]-1], atol=delta/2))[0])
    
    # get subset of x within specified bounds
    x_subset = jax.numpy.array(x[i_start:i_stop])
    y_subset = jax.numpy.array(y[i_start:i_stop])
    
    # return subset of x within specified bounds
    return x_subset, y_subset

# @jax.jit
def master_function(indices) -> tuple:
    # get subset of x and y within bounds params[0] to params[1]
    x_subset, y_subset = isolate_subset(indices)
    
    # call some other functions to do some calculations
    
    # return data subset
    return x_subset, y_subset

# create some arbitrary x and y values
x = np.linspace(0, 10, 1000)
y = np.sin(2*np.pi*x)

# break data down into uniform subsets
subset_indices = np.arange(0, len(x)+1, int(len(x)/5))

# for each subset
for i in range(0, len(subset_indices)-1):
    # get indices of subset
    i_sub = subset_indices[i:i+2]
    
    # analyse subset of data
    x_subset, y_subset = master_function(i_sub)
    
    # create figure
    fig, ax = plt.subplots(tight_layout=True, dpi=200)
    
    # plot original (x, y) and subset of (x, y)
    ax.plot(x, y, "k--", zorder=0, label="original")
    ax.plot(x_subset, y_subset, "r-", zorder=1, label="subset")
    
    # label plot
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.legend()

plt.show()

当给master_function添加JIT装饰后，出现如下错误：

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=0/1)>
The error occurred while tracing the function master_function at untitled0.py:35 for jit. This concrete value was not available in Python because it depends on the value of the argument 'indices'.

我原以为使用jax.numpy替代标准Python/NumPy就不会出现JIT问题，请问如何让isolate_subset()兼容JIT？

解决方案 #

错误原因分析 #

1. 全局变量依赖：isolate_subset直接引用全局numpy数组x和y，JIT编译时无法追踪这些全局变量的依赖关系；同时当Tracer类型的indices尝试索引numpy数组时，会触发Tracer到numpy数组的强制转换，这在JIT追踪流程中是不允许的（Tracer是符号化对象，没有具体数值）。
2. 混合numpy与JAX操作：部分操作直接对numpy数组进行，没有完全迁移到JAX数组体系，导致JIT无法正确追踪计算流程。

修改后的代码 #

import jax
import numpy as np
import matplotlib.pyplot as plt

# enable 64-bit floating point representation in JAX
jax.config.update("jax_enable_x64", True)

def isolate_subset(x: jax.Array, y: jax.Array, indices: jax.Array) -> tuple:
    # compute difference between successive x values
    delta = x[1] - x[0]
    
    # get indices of start and stop of subset
    match_start = jax.numpy.isclose(x, x[indices[0]], atol=delta/2)
    i_start = jax.numpy.min(jax.numpy.where(match_start)[0])
    match_stop = jax.numpy.isclose(x, x[indices[1]-1], atol=delta/2)
    i_stop = jax.numpy.max(jax.numpy.where(match_stop)[0])
    
    # JAX支持动态整数索引的切片操作
    x_subset = x[i_start:i_stop]
    y_subset = y[i_start:i_stop]
    
    return x_subset, y_subset

@jax.jit
def master_function(x: jax.Array, y: jax.Array, indices: jax.Array) -> tuple:
    x_subset, y_subset = isolate_subset(x, y, indices)
    # 这里可以添加后续JIT兼容的计算逻辑
    return x_subset, y_subset

# 创建JAX格式的x和y数组
x = jax.numpy.linspace(0, 10, 1000)
y = jax.numpy.sin(2*jax.numpy.pi*x)

# break data down into uniform subsets
subset_indices = np.arange(0, len(x)+1, int(len(x)/5))

# for each subset
for i in range(0, len(subset_indices)-1):
    # get indices of subset
    i_sub = subset_indices[i:i+2]
    
    # 传入JAX数组进行计算，numpy数组会自动被转换
    x_subset, y_subset = master_function(x, y, i_sub)
    
    # create figure
    fig, ax = plt.subplots(tight_layout=True, dpi=200)
    
    # 转换为numpy数组用于绘图（matplotlib不直接支持JAX数组）
    ax.plot(x, y, "k--", zorder=0, label="original")
    ax.plot(x_subset, y_subset, "r-", zorder=1, label="subset")
    
    # label plot
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.legend()

plt.show()

关键修改点 #

1. 移除全局变量依赖：将x和y作为参数传入isolate_subset和master_function，让JIT能正确追踪所有输入依赖。
2. 使用JAX数组替代numpy数组：将数据集x和y转换为JAX数组，确保所有数组操作都在JAX的可追踪体系内进行。
3. 简化数组转换操作：JAX支持直接对动态整数索引进行切片，无需手动转换为jax.numpy.array，避免额外的类型转换错误。
4. 绘图时转换为numpy数组：由于matplotlib暂不支持直接绘制JAX数组，将结果转换为numpy数组后再绘图（JAX数组可通过np.asarray()或直接索引自动转换）。

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