chore(cc): merge get backend codes #4355
Conversation
Fix deepmodeling#4308. Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
📝 WalkthroughWalkthroughThe changes in this pull request introduce a new function, Changes
Assessment against linked issues
Possibly related PRs
Suggested labels
Suggested reviewers
Warning Rate limit exceeded@njzjz has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 2 minutes and 38 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 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 (3)
source/api_cc/src/DataModifier.cc (1)
Line range hint
31-45: Add missing JAX backend handlingAccording to the PR objectives (issue #4308), JAX backend support should be added. However, the current implementation doesn't include a case for JAX in the backend switch statement.
Add JAX handling similar to other backends. Here's the suggested change:
const DPBackend backend = get_backend(model); if (deepmd::DPBackend::TensorFlow == backend) { #ifdef BUILD_TENSORFLOW dcm = std::make_shared<deepmd::DipoleChargeModifierTF>(model, gpu_rank, name_scope_); #else throw deepmd::deepmd_exception("TensorFlow backend is not built"); #endif } else if (deepmd::DPBackend::PyTorch == backend) { throw deepmd::deepmd_exception("PyTorch backend is not supported yet"); } else if (deepmd::DPBackend::Paddle == backend) { throw deepmd::deepmd_exception("PaddlePaddle backend is not supported yet"); + } else if (deepmd::DPBackend::JAX == backend) { + throw deepmd::deepmd_exception("JAX backend is not supported yet"); } else { throw deepmd::deepmd_exception("Unknown file type"); }source/api_cc/src/DeepSpin.cc (1)
Based on the search results and the previously examined code, I can now provide a final response:
JAX backend implementation missing for DeepSpin
The codebase has a JAX backend implementation for DeepPot (
DeepPotJAX) but lacks a corresponding JAX backend implementation for DeepSpin. This is a potential issue since:
- The DeepSpin class is designed to support multiple backends (as seen in the header file)
- A JAX backend implementation (
DeepSpinJAX) is missing, unlike the PyTorch case referenced in the review comment- The JAX backend would need to implement the
DeepSpinBackendinterface with spin-specific computations🔗 Analysis chain
Line range hint
47-48: Verify JAX backend implementationWhile the code correctly handles the JAX backend case, let's verify if there are any JAX-specific implementations that need to be added similar to the TensorFlow and PyTorch cases.
Let me check the actual JAX implementation to verify if it's properly implemented.
Let me check the JAX backend implementation for DeepSpin specifically.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for JAX backend implementation files # Look for JAX-related header files echo "Searching for JAX header files..." fd -e h -e hpp ".*JAX.*" # Look for JAX-related source files echo "Searching for JAX source files..." fd -e cc -e cpp ".*JAX.*" # Search for JAX-related class definitions echo "Searching for JAX class implementations..." rg -l "class.*JAX"Length of output: 1300
Script:
#!/bin/bash # Check JAX implementation details echo "=== DeepPotJAX header ===" cat source/api_cc/include/DeepPotJAX.h echo -e "\n=== DeepPotJAX implementation ===" cat source/api_cc/src/DeepPotJAX.cc echo -e "\n=== DeepSpin header (for context) ===" cat source/api_cc/include/DeepSpin.hLength of output: 71320
Script:
#!/bin/bash # Check for JAX-specific DeepSpin implementation echo "=== Searching for DeepSpinJAX files ===" fd -e h -e hpp -e cc -e cpp ".*DeepSpinJAX.*" echo -e "\n=== Searching for JAX backend references in DeepSpin implementation ===" rg -l "DeepSpinJAX" source/api_cc/src/DeepSpin.ccLength of output: 345
source/api_cc/src/common.cc (1)
1403-1413: LGTM! Consider adding file validation and case-insensitive extension checks.The implementation correctly determines the backend based on file extensions and aligns with the PR objectives for JAX backend support.
Consider these improvements:
- Add file existence validation before checking extensions
- Make extension checks case-insensitive for better robustness
Example implementation:
deepmd::DPBackend deepmd::get_backend(const std::string& model) { + // Convert to lowercase for case-insensitive comparison + std::string model_lower = model; + std::transform(model_lower.begin(), model_lower.end(), model_lower.begin(), ::tolower); + + // Validate file exists + std::ifstream f(model.c_str()); + if (!f.good()) { + throw deepmd::deepmd_exception("Model file does not exist: " + model); + } + - if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") { + if (model_lower.length() >= 4 && model_lower.substr(model_lower.length() - 4) == ".pth") { return deepmd::DPBackend::PyTorch; - } else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") { + } else if (model_lower.length() >= 3 && model_lower.substr(model_lower.length() - 3) == ".pb") { return deepmd::DPBackend::TensorFlow; - } else if (model.length() >= 11 && model.substr(model.length() - 11) == ".savedmodel") { + } else if (model_lower.length() >= 11 && model_lower.substr(model_lower.length() - 11) == ".savedmodel") { return deepmd::DPBackend::JAX; } throw deepmd::deepmd_exception("Unsupported model file format"); }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (6)
source/api_cc/include/common.h(1 hunks)source/api_cc/src/DataModifier.cc(1 hunks)source/api_cc/src/DeepPot.cc(1 hunks)source/api_cc/src/DeepSpin.cc(1 hunks)source/api_cc/src/DeepTensor.cc(1 hunks)source/api_cc/src/common.cc(1 hunks)
🔇 Additional comments (7)
source/api_cc/src/DataModifier.cc (1)
31-31: Verify get_backend implementation for JAX support
Let's verify that the get_backend function in common.h properly handles JAX model detection.
✅ Verification successful
JAX support is properly implemented in get_backend function
The get_backend implementation in source/api_cc/src/common.cc correctly handles JAX models by checking for the .savedmodel extension and returning DPBackend::JAX. The function properly differentiates between PyTorch (.pth), TensorFlow (.pb), and JAX (.savedmodel) model formats.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for get_backend implementation and its handling of JAX
rg -A 10 "get_backend.*\(" source/api_cc/include/common.h source/api_cc/src/common.cc
Length of output: 1704
source/api_cc/include/common.h (1)
18-23: Documentation could be enhanced with error handling details.
The function declaration and basic documentation look good. However, consider enhancing the documentation with:
- Error handling behavior for unknown/invalid model names
- Expected format of the model name parameter
Example enhancement:
/**
* @brief Get the backend of the model.
* @param[in] model The model name.
* @return The backend of the model.
+ * @details The function analyzes the model name/path to determine the backend.
+ * Returns DPBackend::Unknown if the backend cannot be determined.
+ * @example
+ * get_backend("model.pb") // Returns DPBackend::TensorFlow
+ * get_backend("model.pt") // Returns DPBackend::PyTorch
**/Let's verify the implementation exists and matches this declaration:
✅ Verification successful
Implementation found and matches declaration with clear behavior
The implementation in source/api_cc/src/common.cc shows that:
- Function determines backend based on file extension:
.pth-> PyTorch.pb-> TensorFlow.jax-> JAX (implied by the partial match)
The function is used correctly in multiple places (DeepPot, DeepTensor, DataModifier, DeepSpin) to determine the appropriate backend implementation.
Note: The original review comment suggesting documentation improvements is still valid but not critical, as the implementation clearly shows the expected behavior.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if get_backend implementation exists and matches declaration
# Search for the implementation
echo "Searching for get_backend implementation..."
rg -l "DPBackend\s+get_backend\s*\(.*string.*model"
# Verify the implementation matches the declaration
ast-grep --pattern 'DPBackend get_backend(const std::string& model)'
Length of output: 335
Script:
#!/bin/bash
# Let's search for the implementation in source files
echo "Searching for get_backend implementation in source files..."
rg -A 10 "DPBackend\s+get_backend\s*\(.*string.*model" --type cpp
# Also check if there are any test files that might help understand the behavior
echo -e "\nSearching for related test files..."
rg -l "get_backend.*model" --type cpp "*test*"
Length of output: 1170
Script:
#!/bin/bash
# Let's try a different approach to find the implementation and tests
# First, let's find all cpp files that might contain the implementation
echo "Searching for cpp source files..."
fd -e cpp -e cc
# Then search for get_backend implementation in these files
echo -e "\nSearching for get_backend implementation..."
rg -A 5 "get_backend.*model" --type cpp
# Look for test files with a better pattern
echo -e "\nSearching for test files..."
fd -t f "test" -e cpp -e cc
Length of output: 12669
source/api_cc/src/DeepSpin.cc (1)
Line range hint 39-54: LGTM: Backend handling implementation aligns with PR objectives
The changes effectively implement the backend selection logic by:
- Using the new
get_backendfunction to determine the backend type - Adding explicit support for JAX backend in the conditional checks
- Maintaining proper error handling for unsupported backends
The implementation is consistent with the PR objectives of updating JAX backend handling.
source/api_cc/src/DeepPot.cc (3)
Line range hint 40-65: LGTM! Backend initialization changes look good.
The implementation correctly:
- Uses the new
get_backendfunction for backend determination - Handles JAX backend appropriately through TensorFlow build infrastructure
- Maintains proper error handling and build flags
Line range hint 66-600: No changes in compute methods.
The compute methods and their template instantiations remain unchanged and are not affected by the backend initialization changes.
Line range hint 601-900: Verify JAX model handling in DeepPotModelDevi.
The class should correctly handle JAX models through its DeepPot instances. Let's verify the implementation:
✅ Verification successful
Let me gather more information about how DeepPotModelDevi interacts with DeepPot instances to verify the JAX handling.
Let me try one more search with corrected paths to find the DeepPotModelDevi implementation.
Let me check one final aspect about how DeepPotModelDevi initializes the DeepPot instances.
Based on the evidence gathered from all the searches, I can now provide the final response:
DeepPotModelDevi correctly handles JAX models through delegation
The implementation shows that DeepPotModelDevi acts as a pure delegation layer, passing all computation requests to its underlying DeepPot instances without any backend-specific handling. This design ensures that JAX models are handled correctly since:
- All model-specific logic is encapsulated in the DeepPot instances
- No backend-specific code exists in DeepPotModelDevi
- The compute methods are pure pass-through calls to DeepPot's compute methods
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if DeepPotModelDevi properly initializes JAX models
# Search for DeepPotModelDevi usage with JAX models
rg -l "DeepPotModelDevi.*\.jax" || echo "No direct JAX model usage found"
# Check if there are any special cases for JAX in DeepPotModelDevi
ast-grep --pattern 'class DeepPotModelDevi {
$$$
if $COND {
$$$
jax
$$$
}
$$$
}'
Length of output: 209
Script:
#!/bin/bash
# Check how DeepPotModelDevi initializes and uses DeepPot instances
# Look for DeepPotModelDevi constructor and initialization
ast-grep --pattern 'class DeepPotModelDevi {
$$$
DeepPotModelDevi($$$) {
$$$
}
$$$
}'
# Check how DeepPot instances are created/stored
rg "DeepPot.*dps\[" -A 2
# Check if there's any backend-specific handling
rg "backend|jax|tensorflow|torch" source/api_cc/src/DeepPotModelDevi
Length of output: 272
Script:
#!/bin/bash
# Find the correct path for DeepPotModelDevi files
fd DeepPotModelDevi
# Look for constructor implementation
rg "DeepPotModelDevi::" -A 5
# Check how dps vector is used
rg "dps\[.*\].*compute" -A 3
Length of output: 11184
Script:
#!/bin/bash
# Check DeepPotModelDevi initialization of DeepPot instances
rg "void.*init.*DeepPotModelDevi" -A 15 source/api_cc/src/DeepPot.cc
# Look for DeepPot creation in DeepPotModelDevi
rg "new DeepPot" source/api_cc/src/DeepPot.cc
Length of output: 118
source/api_cc/src/DeepTensor.cc (1)
33-33: LGTM!
The backend detection is now correctly updated to use get_backend(model), allowing for dynamic backend support.
Fix #4308.
Summary by CodeRabbit
New Features
Improvements
Bug Fixes