如何在JAX中对指定函数使用vmap处理批量向量？

我有一个适用于单个向量的函数vec_to_board，代码如下：

def vec_to_board(vector, player, dim, reverse=False):
    player_board = np.zeros(dim * dim)
    player_pos = np.argwhere(vector == player)
    if not reverse:
        player_board[mapping[player_pos.T]] = 1
    else:
        player_board[reverse_mapping[player_pos.T]] = 1
    return np.reshape(player_board, [dim, dim])

现在我希望它能处理批量向量，尝试了以下代码：

states = jnp.array([[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2], [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2]])
batch_size = 1
b_states = vmap(vec_to_board)((states, 1, 4), batch_size)

但这段代码无法正常工作，我理解vmap应该可以完成这类批量转换，请问该如何解决？

问题根源 #

1. numpy与jax数组混用：原函数使用np操作jax数组，会导致兼容性问题，jax的vmap要求操作jax原生数组。
2. vmap参数用法错误：jax的vmap不需要batch_size参数，而是通过in_axes指定每个输入的批量维度；原代码的参数传递方式不符合vmap的要求。
3. 索引处理不兼容批量场景：np.argwhere在批量输入下的维度结构和单向量不同，需要调整索引逻辑适配批量。

修改步骤及代码 #

1. 适配jax的函数版本 #

把原函数中的numpy操作替换为jax.numpy（jnp），并调整索引逻辑：

import jax.numpy as jnp

def vec_to_board(vector, player, dim, reverse=False):
    # 用jnp创建数组，兼容jax的自动微分和vmap
    player_board = jnp.zeros(dim * dim)
    # 用jnp.where获取符合条件的索引，比argwhere更适合批量场景
    player_pos = jnp.where(vector == player)[0]
    # 根据reverse选择映射，这里假设mapping和reverse_mapping是jax数组
    mapped_pos = mapping[player_pos] if not reverse else reverse_mapping[player_pos]
    # 用jnp.at实现原地更新（jax不可变数组需要用at操作）
    player_board = player_board.at[mapped_pos].set(1)
    return jnp.reshape(player_board, [dim, dim])

2. 正确使用vmap批量调用 #

通过in_axes指定批量维度：states的批量维度在第0轴，其他参数（player、dim、reverse）是标量，不需要批量，所以对应None：

from jax import vmap

# 假设mapping和reverse_mapping已定义为jax数组，示例：
mapping = jnp.arange(16)
reverse_mapping = jnp.arange(16)[::-1]

states = jnp.array([
    [1,1,1,0,0,0,0,0,0,0,0,0,0,2,2,2],
    [1,1,1,0,0,0,0,0,0,0,0,0,0,2,2,2]
])

# 创建批量处理函数，指定in_axes：第一个参数（vector）的批量轴是0，其他参数无批量轴
batch_vec_to_board = vmap(vec_to_board, in_axes=(0, None, None, False))
# 调用批量函数
b_states = batch_vec_to_board(states, 1, 4)

关键说明 #

• jax数组不可变性：jax数组是不可变的，不能直接像numpy那样赋值，必须用jnp.at[...]操作实现更新。
• in_axes的作用：in_axes告诉vmap每个输入参数哪个维度是批量维度，None表示该参数在所有批量样本中保持不变。
• 索引逻辑调整：用jnp.where替代np.argwhere，返回的索引结构更适合批量处理，避免转置T带来的维度混乱。

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