2 months ago
- JAX Arrays are immutable, aiding optimized compilation but can limit expressiveness and performance for tasks like memory control.
- Refs are mutable arrays compatible with JAX transformations like jit and grad, allowing in-place reads and writes via indexing.
- Refs have constraints: they cannot be returned from jit functions, passed multiple times, or used in vmap with closed-over refs.
- Automatic differentiation works with pure functions using Refs internally, and ref arguments can be used for gradient accumulation via vjp.with_refs.
- Refs support in-place updates, enabling efficient gradient accumulation over microbatches without materializing separate arrays.
- foreach decorator uses Refs to write flexible loops with non-sequential access, avoiding scan's fixed pattern for better performance.