JAX中jnp.arange使用Abstract value报错及静态参数问题求助

错误原因解析 #

第一个错误（ConcretizationTypeError） #

jnp.arange的stop参数要求编译时可确定的具体数值，但你的size是由动态输入DF[i,j]计算得到的。在JIT编译阶段，这个值是未确定的追踪器（tracer），无法生成固定形状的数组，因此触发报错。

第二个错误（ValueError） #

静态参数（static_argnums指定的参数）必须是可哈希的Python原生类型（如int、float、tuple），但你传入的size是JAX动态计算产生的追踪器对象，这类对象不可哈希，不符合静态参数的要求，导致报错。

解决方案 #

方案1：将size转为Python原生int后作为静态参数传入 #

在调用函数前，先把size计算为Python原生int类型，再传入JIT编译的函数：

from functools import partial
import jax
import jax.numpy as jnp

@partial(jax.jit, static_argnums=3)
def gaussian_shape5(i, j, DF, size):
    """Generate a Gaussian shape."""
    sigma = DF[i, j]
    sigma_sqrt = 2 * (sigma ** 2)
    x = jnp.arange(0, size) - jnp.floor((size - 1) / 2)
    exponent = jnp.exp(-(x ** 2) / sigma_sqrt)
    exponent = (exponent * exponent[:, jnp.newaxis]) / jnp.sum(exponent)
    return exponent

# 调用示例
DF = jnp.array([[1.0, 2.0], [3.0, 4.0]])
i, j = 0, 0

# 先计算size为Python int
sigma_concrete = DF[i, j].item()  # 转为Python float
size_concrete = int(2 * (3 * jnp.ceil(jnp.array(sigma_concrete))) + 1)  # 转为Python int

result = gaussian_shape5(i, j, DF, size_concrete)

方案2：启用JAX动态形状编译（JAX 0.4.13+） #

如果你的JAX版本在0.4.13及以上，可直接启用动态形状支持，无需手动处理静态参数：

import jax
import jax.numpy as jnp

@jax.jit(dynamic=True)
def gaussian_shape5(i, j, DF):
    """Generate a Gaussian shape."""
    sigma = DF[i, j]
    size = jnp.int16(2 * (3 * jnp.ceil(sigma)) + 1)
    sigma_sqrt = 2 * (sigma ** 2)
    x = jnp.arange(0, size) - jnp.floor((size - 1) / 2)
    exponent = jnp.exp(-(x ** 2) / sigma_sqrt)
    exponent = (exponent * exponent[:, jnp.newaxis]) / jnp.sum(exponent)
    return exponent

# 调用示例
DF = jnp.array([[1.0, 2.0], [3.0, 4.0]])
i, j = 0, 0
result = gaussian_shape5(i, j, DF)

方案3：使用jax.lax.dynamic_arange（兼容旧版本JAX） #

若无法升级JAX版本，可用jax.lax.dynamic_arange替代jnp.arange，它支持动态的stop参数：

import jax
import jax.numpy as jnp

@jax.jit
def gaussian_shape5(i, j, DF):
    """Generate a Gaussian shape."""
    sigma = DF[i, j]
    size = jnp.int16(2 * (3 * jnp.ceil(sigma)) + 1)
    sigma_sqrt = 2 * (sigma ** 2)
    # 使用dynamic_arange支持动态stop
    x = jax.lax.dynamic_arange(0, size) - jnp.floor((size - 1) / 2)
    exponent = jnp.exp(-(x ** 2) / sigma_sqrt)
    exponent = (exponent * exponent[:, jnp.newaxis]) / jnp.sum(exponent)
    return exponent

# 调用示例
DF = jnp.array([[1.0, 2.0], [3.0, 4.0]])
i, j = 0, 0
result = gaussian_shape5(i, j, DF)

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