Refs: Mutable Arrays in jax
2 days ago
- JAX数组是不可变的，这有助于优化编译，但对于内存控制等任务可能限制表达能力和性能。
- Refs是可变数组，与JIT和梯度等JAX转换兼容，允许通过索引进行原地读写操作。
- Refs存在限制：无法从JIT函数返回、不能多次传递，也不能在闭包引用Refs时用于vmap转换。
- 自动微分可与内部使用Refs的纯函数配合工作，且可通过vjp.with_refs使用ref参数进行梯度累积。
- Refs支持原地更新，无需实例化独立数组即可在微批次上实现高效的梯度累积。
- foreach装饰器利用Refs编写具有非顺序访问的灵活循环，避免scan的固定模式以提升性能。