scatter 把更新加了上去
主线位置：张量主线的第二步「计算」，按下标写入（JAX 的
.at[i].set）。
版本：JAX 0.12.0.1（CPU，x64）、tenferro-rs 0.8.0（CPU，native）
flowchart TB w["写 99 到下标 1"] --> js["JAX at set"] js --> j99["10, 99, 12"] w --> ts["tenferro scatter"] ts --> t110["10, 110, 12"] w --> fix["掩码加 where"] fix --> j99
实验代码：JAX 脚本一个，tenferro 程序 p05_scatter（traced，CpuEngine::serial()，release）。平台 M2 Max、macOS 27、rustc 1.99.0。
现象
JAX 数组不可变，按下标写入用 .at[...]，set 和 add 是两种操作：
x = jnp.array([10.0, 11.0, 12.0])
x.at[1].set(99.); x.at[1].add(99.); x.at[jnp.array([1, 1])].set(99.); x.at[5].set(99.); x.at[-1].set(99.)
x.at[1].set(99.) [10.0, 99.0, 12.0]
x.at[1].add(99.) [10.0, 110.0, 12.0]
x.at[jnp.array([1, 1])].set(99.) [10.0, 99.0, 12.0]
x.at[5].set(99.) [10.0, 11.0, 12.0]
x.at[-1].set(99.) [10.0, 11.0, 99.0]
tenferro 里对应的原语只有一个 scatter。把 .at[i].set(99) 换成它：
let scfg = ScatterConfig {
update_window_dims: vec![],
inserted_window_dims: vec![0],
scatter_dims_to_operand_dims: vec![0],
index_vector_dim: 1,
};
x.scatter(&idx(&[1])?, &upd99, scfg)? // idx: [n, 1] 的 I64 下标
scatter idx [1] upd 99: [10.0, 110.0, 12.0]
scatter idx [1, 1] upd 99: [10.0, 209.0, 12.0]
scatter idx [5] upd 99: [10.0, 11.0, 111.0]
scatter idx [-1] upd 99: [109.0, 11.0, 12.0]
四种情况和 .at[].set 都不一样：
- 下标 1：加上去，得 110。和
.at[1].add相同。 - 重复下标
[1, 1]：加两次，得 209。 - 越界下标 5：JAX 丢弃这个更新，tenferro 把它夹到最后一个元素加进去，得 111。
- 负下标 −1：JAX 写到最后一个，tenferro 夹到第一个，得 109。
出处
tenferro 的 scatter 是加法 scatter。CPU 内核 scatter_with_pool 拒绝 Bool 时报的就是「Bool data tensors are not supported by additive scatter」；原语规格 primitive-catalog.md 给它的备注是「Transpose of Gather. Needed for AD of Gather and for embed_diag」：它首先是给 gather 求导用的，而 gather 的反向必须把重复下标的余切加在一起。TracedTensor::scatter 的文档写明下标「is clamped rather than rejected」，负下标也不换算（同x[-1] 取到的是第一个元素）。JAX 的 .at[i].set 是另一个 StableHLO scatter（更新函数是「取新值」），越界时默认丢弃。并排对照见整数运算的下标一节。
修法
想要的是 .at[i].add、下标都在范围内时，scatter 就是对的。想要 .at[i].set 时，把 JAX 在下标上做的事搬到主机上（负下标加 n、丢掉越界的、重复的只留最后一个），再用两次 scatter 和 where_select 把「累加」变成「写入」：
fn set_at(x: &TracedTensor, ids: &[i64], upd: &[f64]) -> R<TracedTensor> {
let n = x.concrete_shape()?[0] as i64;
let mut kept: Vec<(i64, f64)> = Vec::new();
for (&i, &u) in ids.iter().zip(upd) {
let i = if i < 0 { i + n } else { i };
if (0..n).contains(&i) {
kept.retain(|(j, _)| *j != i);
kept.push((i, u));
}
}
let zeros = TracedTensor::from_vec_col_major(vec![n as usize], vec![0.0_f64; n as usize])?;
if kept.is_empty() {
return Ok(x.clone());
}
let ki: Vec<i64> = kept.iter().map(|p| p.0).collect();
let ku: Vec<f64> = kept.iter().map(|p| p.1).collect();
let ones = TracedTensor::from_vec_col_major(vec![ki.len()], vec![1.0_f64; ki.len()])?;
let vals = TracedTensor::from_vec_col_major(vec![ki.len()], ku)?;
let mask = zeros.scatter(&idx(&ki)?, &ones, scfg())?.compare(&zeros, CompareDir::Gt)?;
let written = zeros.scatter(&idx(&ki)?, &vals, scfg())?;
Ok(TracedTensor::where_select(&mask, &written, x)?)
}
set_at idx [1] upd 99: [10.0, 99.0, 12.0]
set_at idx [1, 1] upd 99: [10.0, 99.0, 12.0]
set_at idx [5] upd 99: [10.0, 11.0, 12.0]
set_at idx [-1] upd 99: [10.0, 11.0, 99.0]
四行和 JAX 的 .at[].set 逐项相同。往全 0 张量里 scatter 时，下标已经去重，「累加」和「写入」结果一样；掩码决定哪些位置用新值。这个写法要求下标在主机上已知。下标是图里算出来的时，去重和越界检查没有现成的原语，要按具体问题另想办法（例如保证下标唯一且在范围内，再用这里的掩码法）。
怎么确认没中招
- 在
scatter的测试里让原值不为 0（例如 10、11、12）：原值是 0 时累加和写入看不出区别。 - 测一次重复下标和一次越界下标，期望值按你要的 JAX 操作（
set还是add）写。 - 代码评审时把每个
scatter调用旁边注明它对应.at[].add还是.at[].set。
小结
- tenferro 的
scatter是累加，对应.at[i].add，不是.at[i].set；重复下标加多次。 - 越界和负下标都夹到边上再加进去，JAX 的
set/add是丢弃越界、负下标从末尾数。 - 要
set：主机上整理下标，往全 0 张量里 scatter 出新值和掩码，用where_select合成。