Skip to content

fix(jax): handle DPA-2 pbc/nopbc without mapping #4363

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Nov 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions deepmd/jax/jax2tf/serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -294,6 +294,12 @@
)

tf_model.get_model_def_script = get_model_def_script

@tf.function
def has_message_passing() -> tf.Tensor:
return tf.constant(model.has_message_passing(), dtype=tf.bool)

Check warning on line 300 in deepmd/jax/jax2tf/serialization.py

View check run for this annotation

Codecov / codecov/patch

deepmd/jax/jax2tf/serialization.py#L300

Added line #L300 was not covered by tests

tf_model.has_message_passing = has_message_passing
tf.saved_model.save(
tf_model,
model_file,
Expand Down
2 changes: 2 additions & 0 deletions source/api_cc/include/DeepPotJAX.h
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,8 @@ class DeepPotJAX : public DeepPotBackend {
std::vector<int64_t> sel;
// number of neighbors
int nnei;
// do message passing
bool do_message_passing;
// padding to nall
int padding_to_nall = 0;
// padding for nloc
Expand Down
24 changes: 23 additions & 1 deletion source/api_cc/src/DeepPotJAX.cc
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,12 @@
return TF_INT64;
}

struct tf_function_not_found : public deepmd::deepmd_exception {
public:
tf_function_not_found() : deepmd_exception() {};
tf_function_not_found(const std::string& msg) : deepmd_exception(msg) {};
};

inline TFE_Op* get_func_op(TFE_Context* ctx,
const std::string func_name,
const std::vector<TF_Function*>& funcs,
Expand All @@ -72,7 +78,7 @@
TF_Function* func = NULL;
find_function(func, funcs, func_name);
if (func == NULL) {
throw std::runtime_error("Function " + func_name + " not found");
throw tf_function_not_found("Function " + func_name + " not found");
}
const char* real_func_name = TF_FunctionName(func);
// execute the function
Expand Down Expand Up @@ -314,6 +320,13 @@
ntypes = type_map_.size();
sel = get_vector<int64_t>(ctx, "get_sel", func_vector, device, status);
nnei = std::accumulate(sel.begin(), sel.end(), decltype(sel)::value_type(0));
try {
do_message_passing = get_scalar<bool>(ctx, "do_message_passing",
func_vector, device, status);
} catch (tf_function_not_found& e) {
// compatibile with models generated by v3.0.0rc0
do_message_passing = false;
}
inited = true;
}

Expand Down Expand Up @@ -584,6 +597,15 @@
for (size_t ii = 0; ii < nall_real; ii++) {
mapping[ii] = lmp_list.mapping[fwd_map[ii]];
}
} else if (nloc_real == nall_real) {
// no ghost atoms
for (size_t ii = 0; ii < nall_real; ii++) {
mapping[ii] = ii;
}
} else if (do_message_passing) {
throw deepmd::deepmd_exception(

Check warning on line 606 in source/api_cc/src/DeepPotJAX.cc

View check run for this annotation

Codecov / codecov/patch

source/api_cc/src/DeepPotJAX.cc#L606

Added line #L606 was not covered by tests
"Mapping is required for a message passing model. If you are using "
"LAMMPS, set `atom_modify map yes`");
}
input_list[3] = add_input(op, mapping, mapping_shape, data_tensor[3], status);
// fparam
Expand Down
Loading