如何在lax.while_loop中使用传递变量实现dynamic_slice？

在lax.while_loop()中使用传递的变量执行dynamic_slice动态切片时触发错误，已知简单函数中可通过partial传递静态值，但不清楚如何处理作为传递值的变量（此处为length）。

报错信息 #

new_u = lax.dynamic_slice(u,(0,0),(0,length-1))
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: Shapes must be 1D sequences of concrete values of integer type, got (0, Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=2/0)>).
If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.

示例代码 #

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

u = jnp.array([[1,2,3,4,5],[0,0,0,0,0]])

def body_fun(carry):
    length, sum_u = carry
    new_u = lax.dynamic_slice(u,(0,0),(0,length-1))
    #new_u = lax.dynamic_slice(u,(0,0),(2,4))
    jax.debug.print("new_u:{}", new_u)
    new_sum_u = jnp.sum(new_u)
    new_length = length -1
    return (new_length, new_sum_u)

def cond_fun(carry):
    length, sum_u = carry
    keep_condition = sum_u < 5
    return keep_condition

init_carry = (5,10)
out = lax.while_loop(cond_fun, body_fun, init_carry)
print(out)

问题根源 #

1. 无效的切片维度：你在dynamic_slice中指定的切片形状为(0, length-1)，但原数组u的第一维度是2，第一维度设为0会生成空数组，同时违背了数组切片的基本逻辑。
2. 动态变量类型模糊：length是while_loop传递的动态变量，JAX追踪时将其识别为弱类型值，而dynamic_slice要求形状参数必须是明确的整数类型。

修正代码 #

import jax.numpy as jnp
from jax import lax
import jax

u = jnp.array([[1,2,3,4,5],[0,0,0,0,0]])

def body_fun(carry):
    length, sum_u = carry
    # 修正第一维度为原数组的维度大小，同时明确动态变量类型
    slice_shape = (2, jax.lax.convert_element_type(length - 1, jnp.int32))
    new_u = lax.dynamic_slice(u, (0, 0), slice_shape)
    jax.debug.print("new_u:{}", new_u)
    new_sum_u = jnp.sum(new_u)
    new_length = length - 1
    return (new_length, new_sum_u)

def cond_fun(carry):
    length, sum_u = carry
    return sum_u < 5

init_carry = (5, 10)
out = lax.while_loop(cond_fun, body_fun, init_carry)
print(out)

核心要点 #

• 正确设置切片形状：切片形状的各维度必须与原数组对应维度匹配或更小，这里将第一维度设为2（原数组第一维度的大小），才能切取出有效子数组。
• 明确动态变量类型：使用jax.lax.convert_element_type将动态计算的length-1转换为int32类型，消除JAX追踪时的弱类型歧义，满足dynamic_slice的参数要求。
• while_loop动态变量处理：while_loop内部的变量天然支持动态追踪，无需借助partial传递静态值，只要确保所有操作符合JAX的动态计算规则即可——dynamic_slice本身支持动态形状参数，只需保证参数类型正确。

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