如何在一行代码中同时修改Jax数组的多个维度？

你遇到的_IndexUpdateRef' object is not subscriptable错误，本质是jax_array.at[...]返回的是更新引用对象，没法再链式做下标操作。要在一行代码里实现需求，直接在at中同时指定行和列的多维索引就行，不用拆分两次下标。

举个具体的实现例子：

import jax
import jax.numpy as jnp

# 初始化5x5的JAX数组
jax_array = jnp.ones((5, 5))
# 准备要替换的2x2数组
update_data = jnp.array([[10, 20], [30, 40]])
# 指定目标行和列
target_rows = [0, 3]
target_cols = [1, 2]

# 一行完成更新
updated_array = jax_array.at[target_rows, target_cols].set(update_data)

这里直接在at里用逗号分隔行、列索引，JAX会自动匹配出2x2的目标区域，和你要替换的数组形状对应，完美完成赋值。如果你想更明确地标注多维索引，也可以写成at[(target_rows, target_cols)]，效果完全一致。

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