docs: document JAX backend #4259
Conversation
Signed-off-by: Jinzhe Zeng <[email protected]>
Signed-off-by: Jinzhe Zeng <[email protected]>
📝 WalkthroughWalkthroughThe changes in this pull request primarily involve the addition of support for the JAX backend in the DeePMD-kit documentation. This includes updates to various documents to specify JAX-related filename extensions, environment variables, installation instructions, and command usage. The documentation now reflects JAX as a supported backend alongside existing backends like TensorFlow and PyTorch, while retaining the content related to those backends unchanged. Changes
Possibly related PRs
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🔇 Additional comments (2)
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (10)
doc/model/sel.md (1)
27-34: LGTM! Consider adding a newline for consistency.The JAX backend command is correctly documented and follows the same pattern as TensorFlow and PyTorch backends. The command structure and parameters are consistent with other backends.
For consistency with other tabs, consider removing the extra newline after the JAX tab-item declaration:
:::{tab-item} JAX {{ jax_icon }} - ```sh dp --jax neighbor-stat -s data -r 6.0 -t O H
:::
</blockquote></details> <details> <summary>doc/backend.md (3)</summary><blockquote> Line range hint `3-6`: **Update the introduction to include JAX backend.** The introduction paragraph needs to be updated to include JAX as a supported backend. Apply this diff: ```diff -DeePMD-kit supports multiple backends: TensorFlow and PyTorch. +DeePMD-kit supports multiple backends: TensorFlow, PyTorch, and JAX. To use DeePMD-kit, you must install at least one backend. Each backend does not support all features. -In the documentation, TensorFlow {{ tensorflow_icon }} and PyTorch {{ pytorch_icon }} icons are used to mark whether a backend supports a feature. +In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, and JAX {{ jax_icon }} icons are used to mark whether a backend supports a feature.
Line range hint
73-74: Document JAX backend switching commands.The "Training" subsection under "Switch the backend" should mention that JAX doesn't support training, as stated in the JAX backend section.
Apply this diff:
When training and freezing a model, you can use `dp --tf` or `dp --pt` in the command line to switch the backend. +Note that the JAX backend currently does not support training.
Line range hint
78-83: Document JAX backend conversion compatibility.The backend conversion section should clarify JAX's compatibility status with other backends.
Apply this diff after the warning block:
Currently, only the `se_e2_a` model fully supports the backend conversion between TensorFlow {{ tensorflow_icon }} and PyTorch {{ pytorch_icon }}. +The JAX backend currently does not support model conversion to or from other backends. :::doc/env.md (1)
34-34: Consider documenting additional JAX environment variables.The documentation correctly includes the two most commonly used JAX environment variables with proper links. However, consider adding these additional important JAX environment variables that users might need:
CUDA_VISIBLE_DEVICES: While already mentioned in CUDA section, it's particularly important for JAXXLA_PYTHON_CLIENT_MEM_FRACTION: Controls memory allocationJAX_ENABLE_X64: Controls double precision supportHere's a suggested expansion:
- - {{ jax_icon }} [`JAX_PLATFORMS`](https://jax.readthedocs.io/en/latest/faq.html#controlling-data-and-computation-placement-on-devices) and [`XLA_FLAGS`](https://jax.readthedocs.io/en/latest/gpu_performance_tips.html#xla-performance-flags) are commonly used. + - {{ jax_icon }} Commonly used JAX environment variables: + - [`JAX_PLATFORMS`](https://jax.readthedocs.io/en/latest/faq.html#controlling-data-and-computation-placement-on-devices): Controls which platforms (e.g., 'cpu', 'gpu') JAX should initialize + - [`XLA_FLAGS`](https://jax.readthedocs.io/en/latest/gpu_performance_tips.html#xla-performance-flags): Controls XLA compiler behavior + - `XLA_PYTHON_CLIENT_MEM_FRACTION`: Controls the fraction of GPU memory allocated to JAX + - `JAX_ENABLE_X64`: Enables double precision support (disabled by default)doc/model/train-fitting-dos.md (2)
Line range hint
1-200: Add JAX command examples for training and testing.The document shows command examples for TensorFlow and PyTorch but is missing the equivalent JAX commands. Please add the JAX tab with appropriate commands in both the training and testing sections.
Add these tab sections:
::::{tab-set} :::{tab-item} TensorFlow {{ tensorflow_icon }} ... ::: :::{tab-item} PyTorch {{ pytorch_icon }} ... ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +dp --jax train input.json +``` + +::: ::::And in the testing section:
::::{tab-set} :::{tab-item} TensorFlow {{ tensorflow_icon }} ... ::: :::{tab-item} PyTorch {{ pytorch_icon }} ... ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +dp --jax freeze -o frozen_model.jax + +dp --jax test -m frozen_model.jax -s ../data/111/$k -d ${output_prefix} -a -n 100 +``` + +::: ::::🧰 Tools
🪛 LanguageTool
[uncategorized] ~6-~6: Possible missing comma found.
Context: ...jax_icon }}, DP {{ dpmodel_icon }} ::: Here we present an API to DeepDOS model, whi...(AI_HYDRA_LEO_MISSING_COMMA)
Based on the gathered information, I can now generate the final response:
Documentation needs JAX-related command examples
The JAX backend is correctly listed in the "Supported backends" section, but the command examples in the tab-set are incomplete:
- The document only shows TensorFlow and PyTorch examples
- JAX command examples (
dp --jax train input.json) are missing in the tab-set- This is inconsistent with other documentation files that list JAX as a supported backend
🔗 Analysis chain
Line range hint
1-200: Verify consistency with other documentation files.Let's ensure the JAX-related information is consistent across the documentation.
Let me gather more information about JAX-related documentation to ensure complete consistency.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check consistency of JAX-related documentation # Test 1: Check file extensions mentioned in backend documentation echo "Checking file extensions..." rg -i "\.jax|\.xlo" doc/backend.md # Test 2: Check environment variables echo "Checking environment variables..." rg -i "JAX_PLATFORMS|XLA_FLAGS" doc/env.md # Test 3: Check installation instructions echo "Checking installation instructions..." rg -i "jax.*pip install" doc/install/install-from-source.md # Test 4: Check command examples in other model docs echo "Checking JAX commands in other docs..." rg -i "dp --jax" doc/model/Length of output: 1011
Script:
#!/bin/bash # Check for JAX-related content in other training documentation files # Look for tab-sets containing backend examples echo "Checking for backend tab-sets in training docs..." rg -A 5 "tab-set" doc/model/train-*.md # Check for backend icons in other docs echo "Checking backend icons usage..." rg "{{ (tensorflow|pytorch|jax|dpmodel)_icon }}" doc/model/train-*.md # Check for backend specifications echo "Checking backend specifications..." rg -i "supported.*backend" doc/model/train-*.mdLength of output: 4614
🧰 Tools
🪛 LanguageTool
[uncategorized] ~6-~6: Possible missing comma found.
Context: ...jax_icon }}, DP {{ dpmodel_icon }} ::: Here we present an API to DeepDOS model, whi...(AI_HYDRA_LEO_MISSING_COMMA)
doc/model/train-se-atten.md (1)
1-1: LGTM! Consider adding JAX version requirement.The addition of JAX backend support is well-documented and follows the existing format. However, consider adding a note about the minimum required JAX version (0.4.33 or above) as mentioned in other documentation files.
Also applies to: 4-4
doc/install/install-from-source.md (2)
81-94: Consider adding version requirements and verification steps.While the installation instructions are clear, consider adding:
- Minimum supported JAX version requirement
- A verification step similar to TensorFlow's tab (e.g., simple JAX code to verify installation)
- Note about CUDA compatibility requirements for GPU support
Example verification step to add:
One can also [use conda](https://docs.deepmodeling.org/faq/conda.html) to install JAX from [conda-forge](https://conda-forge.org). +To verify the installation, run + +```python +import jax +import jax.numpy as jnp +print(jax.device_count()) +print(jnp.sum(jnp.array([1, 2, 3]))) +``` +
81-94: Document JAX-specific configuration options.Consider adding:
- JAX-specific environment variables (if any)
- CMake configuration for JAX backend (e.g.,
ENABLE_JAXflag) in the C++ interface sectionExample CMake variable to document:
:::{cmake:variable} ENABLE_PYTORCH **Type**: `BOOL` (`ON`/`OFF`), Default: `OFF` {{ pytorch_icon }} Whether building the PyTorch backend. ::: + +:::{cmake:variable} ENABLE_JAX +**Type**: `BOOL` (`ON`/`OFF`), Default: `OFF` +{{ jax_icon }} Whether building the JAX backend. +:::
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
⛔ Files ignored due to path filters (1)
doc/_static/jax.svgis excluded by!**/*.svg
📒 Files selected for processing (10)
- doc/backend.md (1 hunks)
- doc/conf.py (1 hunks)
- doc/env.md (1 hunks)
- doc/install/install-from-source.md (1 hunks)
- doc/model/sel.md (1 hunks)
- doc/model/train-energy.md (1 hunks)
- doc/model/train-fitting-dos.md (1 hunks)
- doc/model/train-se-atten.md (1 hunks)
- doc/model/train-se-e2-a.md (1 hunks)
- doc/model/train-se-e2-r.md (1 hunks)
🔇 Additional comments (13)
doc/backend.md (1)
26-33: LGTM! Clear and concise documentation of the JAX backend.The JAX backend section effectively communicates:
- File extensions and their custom nature
- Version requirements
- Current limitations and development status
doc/model/train-se-e2-r.md (2)
1-1: LGTM: JAX backend icon correctly addedThe JAX backend icon is properly added to the descriptor header, maintaining consistency with other supported backends.
4-4: LGTM: JAX backend support properly documentedThe supported backends note is correctly updated to include JAX with its corresponding icon.
doc/env.md (1)
34-34: Verify documentation links.Let's ensure the JAX documentation links are valid and accessible.
✅ Verification successful
Documentation links are valid and accessible
Both JAX documentation links return HTTP 200 status codes, confirming they are active and accessible:
JAX_PLATFORMSdocumentation at/faq.html#controlling-data-and-computation-placement-on-devicesXLA_FLAGSdocumentation at/gpu_performance_tips.html#xla-performance-flags🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify that the JAX documentation links are valid # Expected results: Both URLs should return HTTP 200 status urls=( "https://jax.readthedocs.io/en/latest/faq.html#controlling-data-and-computation-placement-on-devices" "https://jax.readthedocs.io/en/latest/gpu_performance_tips.html#xla-performance-flags" ) for url in "${urls[@]}"; do echo "Checking $url" curl -s -L -o /dev/null -w "%{http_code}" "$url" echo doneLength of output: 968
doc/conf.py (2)
170-170: LGTM! The JAX icon addition follows the existing pattern.The implementation is consistent with other backend icons and uses the same styling class.
170-170: Verify JAX icon file exists.Please ensure that the JAX SVG icon file exists at
_static/jax.svg.✅ Verification successful
JAX icon file is present at the expected location
The verification confirms that the JAX SVG icon file exists at
doc/_static/jax.svg, ensuring the documentation will render correctly with the icon.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check if the JAX icon file exists in the static directory if [ -f "doc/_static/jax.svg" ]; then echo "JAX icon file exists" else echo "JAX icon file is missing" fiLength of output: 83
doc/model/train-fitting-dos.md (1)
1-1: LGTM: JAX backend support added correctly.The addition of JAX backend to the title and supported backends list is consistent with the PR objectives.
Also applies to: 4-4
doc/model/train-se-e2-a.md (3)
1-1: LGTM: JAX backend icon addition is correct.The addition of the JAX icon to the descriptor title follows the existing pattern and correctly indicates JAX backend support.
4-4: LGTM: JAX backend support documentation is accurate.The addition of JAX to the supported backends note is properly formatted and consistent with the documentation style.
1-4: Consider adding JAX-specific usage notes if applicable.While the JAX backend support is properly documented, consider mentioning if there are any JAX-specific considerations, limitations, or configuration options when using this descriptor with the JAX backend.
Let me check if other descriptor documentation files include JAX-specific details:
doc/model/train-energy.md (3)
1-1: LGTM: JAX backend added to title.The addition of the JAX icon to the title is consistent with documenting JAX as a supported backend.
4-4: LGTM: JAX backend added to supported backends list.The addition of JAX to the list of supported backends in the note section is accurate and consistent with the documentation updates.
1-4: Verify documentation consistency.Let's ensure the JAX-related changes are consistent across the documentation:
✅ Verification successful
Documentation consistency verified
The JAX icon is properly defined in
doc/conf.pyand is consistently used across multiple documentation files. The backend support documentation follows a consistent pattern across various files, including proper usage of the JAX icon alongside other backend icons.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Verify JAX icon definition and consistent backend support documentation # Test 1: Check if JAX icon is properly defined echo "Checking JAX icon definition..." rg "jax_icon" # Test 2: Check if other similar documentation files have been updated echo "Checking other documentation files for backend support..." rg -l "Supported backends.*TensorFlow.*PyTorch" doc/Length of output: 2556
Signed-off-by: Jinzhe Zeng <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Outside diff range and nitpick comments (2)
doc/install/easy-install.md (2)
143-146: Fix grammatical error in the extras description.Change "to extras" to "as extras" for better grammatical correctness.
-[The LAMMPS module](../third-party/lammps-command.md) and [the i-PI driver](../third-party/ipi.md) are only provided on Linux and macOS for the TensorFlow backend. To install LAMMPS and/or i-PI, add `lmp` and/or `ipi` to extras: +[The LAMMPS module](../third-party/lammps-command.md) and [the i-PI driver](../third-party/ipi.md) are only provided on Linux and macOS for the TensorFlow backend. To install LAMMPS and/or i-PI, add `lmp` and/or `ipi` as extras:🧰 Tools
🪛 LanguageTool
[uncategorized] ~143-~143: The preposition ‘as’ seems more likely in this position.
Context: ...MPS and/or i-PI, addlmpand/oripito extras: ```bash pip install deepmd-kit...(AI_HYDRA_LEO_REPLACE_TO_AS)
216-216: Improve warning message readability.The current warning message repeats "you want to". Consider making it more concise.
-If your platform is not supported, or you want to build against the installed backends, or you want to enable ROCM support, please [build from source](install-from-source.md). +If your platform is not supported, you need to build against installed backends, or require ROCM support, please [build from source](install-from-source.md).🧰 Tools
🪛 LanguageTool
[style] ~216-~216: You have already used this phrasing in nearby sentences. Consider replacing it to add variety to your writing.
Context: ... against the installed backends, or you want to enable ROCM support, please [build from...(REP_WANT_TO_VB)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (2)
- doc/install/easy-install-dev.md (1 hunks)
- doc/install/easy-install.md (1 hunks)
🧰 Additional context used
🪛 LanguageTool
doc/install/easy-install.md
[uncategorized] ~143-~143: The preposition ‘as’ seems more likely in this position.
Context: ...MPS and/or i-PI, addlmpand/oripito extras: ```bash pip install deepmd-kit...(AI_HYDRA_LEO_REPLACE_TO_AS)
[style] ~216-~216: You have already used this phrasing in nearby sentences. Consider replacing it to add variety to your writing.
Context: ... against the installed backends, or you want to enable ROCM support, please [build from...(REP_WANT_TO_VB)
🔇 Additional comments (2)
doc/install/easy-install-dev.md (1)
Line range hint
25-28: LGTM! Clear backend support specification.The note clearly indicates that the C Library supports only the TensorFlow backend, which is important information for users.
doc/install/easy-install.md (1)
107-108: LGTM: Good addition of environment setup guidance.Adding a link to environment creation documentation is helpful for users to properly set up their workspace before installation.
|
I feel that JAX icons are very different from others...
|
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## devel #4259 +/- ##
==========================================
+ Coverage 84.23% 84.34% +0.11%
==========================================
Files 549 549
Lines 51462 51463 +1
Branches 3051 3051
==========================================
+ Hits 43348 43407 +59
+ Misses 7153 7098 -55
+ Partials 961 958 -3 ☔ View full report in Codecov by Sentry. |
Signed-off-by: Jinzhe Zeng <[email protected]>
Summary by CodeRabbit
Release Notes
New Features
Documentation Updates