Upgrade JAX to v0.8.2 and clean up installation instructions#1
Upgrade JAX to v0.8.2 and clean up installation instructions#1google-labs-jules[bot] wants to merge 2 commits intojax-upgradefrom
Conversation
- Updated JAX to v0.8.2, Flax to v0.12.2, Optax to v0.2.6, and TensorFlow to v2.20.0 in `requirements.txt`. - Removed `jaxlib` from `requirements.txt` to allow flexible installation (CPU vs GPU). - Replaced deprecated `jax.tree_map` with `jax.tree.map` in `calcil/reconstruction.py`. - Updated `README.md` with clearer installation instructions, emphasizing that default installation is CPU-only and providing commands for GPU support. - Fixed an issue with `orbax.checkpoint` requiring absolute paths by using `os.path.abspath` in `calcil/reconstruction.py`.
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
- Updated JAX to v0.8.2, Flax to v0.12.2, Optax to v0.2.6, and TensorFlow to v2.20.0 in `requirements.txt`. - Removed `jaxlib` from `requirements.txt` to allow flexible installation (CPU vs GPU). - Replaced deprecated `jax.tree_map` with `jax.tree.map` in `calcil/reconstruction.py`. - Updated `README.md` with clearer installation instructions, emphasizing that default installation is CPU-only and providing commands for GPU support. - Fixed an issue with `orbax.checkpoint` requiring absolute paths by using `os.path.abspath` in `calcil/reconstruction.py`. - Updated documentation in `docs/` to reflect changes in `requirements.txt` and `README.md`.
This PR upgrades the JAX dependency to v0.8.2 as requested. It also updates compatible versions for Flax, Optax, and TensorFlow. The code was modified to use
jax.tree.mapinstead of the deprecatedjax.tree_map. Installation instructions inREADME.mdwere updated to reflect the new dependency management strategy (CPU by default, explicit instruction for GPU). Verified the upgrade by running a deconvolution example script.PR created automatically by Jules for task 5519985400734555835 started by @rmcao