JAX: commitment issues
14 hours ago
- Creating an array with `jax.default_device` does not commit it to that device; JAX may move it elsewhere.
- Simple array lookups (e.g., `array[0]`) on an uncommitted CPU array can take over 5 seconds due to implicit GPU transfers.
- Explicitly committing an array with `jax.device_put(array, cpu0)` reduces lookup times dramatically (e.g., from 1.2s to <0.0002s).
- The issue stems from JAX moving the entire uncommitted array to the GPU for each lookup, causing severe performance degradation.
- For CPU-bound data like training datasets, always use `device_put` to nail arrays to the desired device and avoid hidden GPU overhead.