fix(jax): skip padding for non-XLA SavedModels

source/api_cc/include/DeepPotJAX.h

@@ -201,6 +201,9 @@ class DeepPotJAX : public DeepPotBackend {
   bool do_message_passing;
   // has default fparam
   bool has_default_fparam_;
+  // whether SavedModel execution goes through XLA and benefits from shape
+  // padding; true for JAX/jax2tf XlaCallModule and TF2 jit_compile exports
+  bool uses_xla_compilation_ = false;
   // padding to nall
   int padding_to_nall = 0;
   // padding for nloc

source/api_cc/src/DeepPotJAX.cc

@@ -9,10 +9,12 @@
 #include <array>
 #include <cstdint>
 #include <cstdio>
+#include <cstring>
 #include <iostream>
 #include <numeric>
 #include <ostream>
 #include <stdexcept>
+#include <string>
 #include <vector>
 
 #include "common.h"
@@ -48,6 +50,265 @@ inline void find_function(TF_Function*& found_func,
   found_func = NULL;
 }
 
+inline bool operation_attr_bool_true(TF_Operation* oper,
+                                     const char* attr_name) {
+  TF_Status* attr_status = TF_NewStatus();
+  unsigned char value = 0;
+  TF_OperationGetAttrBool(oper, attr_name, &value, attr_status);
+  const bool result = TF_GetCode(attr_status) == TF_OK && value != 0;
+  TF_DeleteStatus(attr_status);
+  return result;
+}
+
+inline bool read_proto_varint(const char*& ptr,
+                              const char* end,
+                              uint64_t& value) {
+  value = 0;
+  int shift = 0;
+  while (ptr < end && shift <= 63) {
+    const uint64_t byte = static_cast<unsigned char>(*ptr++);
+    value |= (byte & 0x7f) << shift;
+    if ((byte & 0x80) == 0) {
+      return true;
+    }
+    shift += 7;
+  }
+  return false;
+}
+
+inline bool read_proto_bytes(const char*& ptr,
+                             const char* end,
+                             const char*& payload,
+                             size_t& payload_size) {
+  uint64_t size = 0;
+  if (!read_proto_varint(ptr, end, size) ||
+      size > static_cast<uint64_t>(end - ptr)) {
+    return false;
+  }
+  payload = ptr;
+  payload_size = static_cast<size_t>(size);
+  ptr += payload_size;
+  return true;
+}
+
+inline bool skip_proto_field(const int wire_type,
+                             const char*& ptr,
+                             const char* end) {
+  uint64_t ignored = 0;
+  const char* payload = NULL;
+  size_t payload_size = 0;
+  switch (wire_type) {
+    case 0:
+      return read_proto_varint(ptr, end, ignored);
+    case 1:
+      if (end - ptr < 8) {
+        return false;
+      }
+      ptr += 8;
+      return true;
+    case 2:
+      return read_proto_bytes(ptr, end, payload, payload_size);
+    case 5:
+      if (end - ptr < 4) {
+        return false;
+      }
+      ptr += 4;
+      return true;
+    default:
+      return false;
+  }
+}
+
+inline bool proto_bytes_equal(const char* payload,
+                              const size_t payload_size,
+                              const std::string& expected) {
+  return payload_size == expected.size() &&
+         std::memcmp(payload, expected.data(), payload_size) == 0;
+}
+
+inline bool attr_value_bool_true(const char* data, const size_t size) {
+  const char* ptr = data;
+  const char* end = data + size;
+  while (ptr < end) {
+    uint64_t tag = 0;
+    if (!read_proto_varint(ptr, end, tag)) {
+      return false;
+    }
+    const int field_number = static_cast<int>(tag >> 3);
+    const int wire_type = static_cast<int>(tag & 0x7);
+    // AttrValue.b = 5, encoded as a varint bool.
+    if (field_number == 5 && wire_type == 0) {
+      uint64_t value = 0;
+      return read_proto_varint(ptr, end, value) && value != 0;
+    }
+    if (!skip_proto_field(wire_type, ptr, end)) {
+      return false;
+    }
+  }
+  return false;
+}
+
+inline bool function_attr_bool_true(TF_Function* func, const char* attr_name) {
+  TF_Status* attr_status = TF_NewStatus();
+  TF_Buffer* attr_value = TF_NewBuffer();
+  TF_FunctionGetAttrValueProto(func, attr_name, attr_value, attr_status);
+  const bool result =
+      TF_GetCode(attr_status) == TF_OK && attr_value->data != NULL &&
+      attr_value_bool_true(static_cast<const char*>(attr_value->data),
+                           attr_value->length);
+  TF_DeleteBuffer(attr_value);
+  TF_DeleteStatus(attr_status);
+  return result;
+}
+
+inline bool attr_entry_is_xla_must_compile_true(const char* data,
+                                                const size_t size) {
+  const char* ptr = data;
+  const char* end = data + size;
+  bool key_matches = false;
+  bool value_is_true = false;
+  while (ptr < end) {
+    uint64_t tag = 0;
+    if (!read_proto_varint(ptr, end, tag)) {
+      return false;
+    }
+    const int field_number = static_cast<int>(tag >> 3);
+    const int wire_type = static_cast<int>(tag & 0x7);
+    if (wire_type == 2 && (field_number == 1 || field_number == 2)) {
+      const char* payload = NULL;
+      size_t payload_size = 0;
+      if (!read_proto_bytes(ptr, end, payload, payload_size)) {
+        return false;
+      }
+      if (field_number == 1) {
+        key_matches =
+            proto_bytes_equal(payload, payload_size, "_XlaMustCompile");
+      } else if (field_number == 2) {
+        value_is_true = attr_value_bool_true(payload, payload_size);
+      }
+    } else if (!skip_proto_field(wire_type, ptr, end)) {
+      return false;
+    }
+  }
+  return key_matches && value_is_true;
+}
+
+inline bool node_def_uses_xla(const char* data, const size_t size) {
+  const char* ptr = data;
+  const char* end = data + size;
+  while (ptr < end) {
+    uint64_t tag = 0;
+    if (!read_proto_varint(ptr, end, tag)) {
+      return false;
+    }
+    const int field_number = static_cast<int>(tag >> 3);
+    const int wire_type = static_cast<int>(tag & 0x7);
+    if (wire_type == 2 && (field_number == 2 || field_number == 5)) {
+      const char* payload = NULL;
+      size_t payload_size = 0;
+      if (!read_proto_bytes(ptr, end, payload, payload_size)) {
+        return false;
+      }
+      // NodeDef.op = 2. This identifies jax2tf native serialization.
+      if (field_number == 2 &&
+          proto_bytes_equal(payload, payload_size, "XlaCallModule")) {
+        return true;
+      }
+      // NodeDef.attr = 5. This catches PartitionedCall nodes marked by
+      // tf.function(jit_compile=True).
+      if (field_number == 5 &&
+          attr_entry_is_xla_must_compile_true(payload, payload_size)) {
+        return true;
+      }
+    } else if (!skip_proto_field(wire_type, ptr, end)) {
+      return false;
+    }
+  }
+  return false;
+}
+
+inline bool function_def_uses_xla(const char* data, const size_t size) {
+  // TensorFlow's C API exposes TF_Function bodies only as serialized
+  // FunctionDef protos. Use a minimal wire-format reader over the stable
+  // FunctionDef/NodeDef/AttrValue field numbers instead of a raw byte
+  // substring, and avoid depending on TensorFlow C++ protobuf headers.
+  const char* ptr = data;
+  const char* end = data + size;
+  while (ptr < end) {
+    uint64_t tag = 0;
+    if (!read_proto_varint(ptr, end, tag)) {
+      return false;
+    }
+    const int field_number = static_cast<int>(tag >> 3);
+    const int wire_type = static_cast<int>(tag & 0x7);
+    if (wire_type == 2 && (field_number == 3 || field_number == 5)) {
+      const char* payload = NULL;
+      size_t payload_size = 0;
+      if (!read_proto_bytes(ptr, end, payload, payload_size)) {
+        return false;
+      }
+      // FunctionDef.node_def = 3.
+      if (field_number == 3 && node_def_uses_xla(payload, payload_size)) {
+        return true;
+      }
+      // FunctionDef.attr = 5. This catches concrete functions marked by
+      // tf.function(jit_compile=True).
+      if (field_number == 5 &&
+          attr_entry_is_xla_must_compile_true(payload, payload_size)) {
+        return true;
+      }
+    } else if (!skip_proto_field(wire_type, ptr, end)) {
+      return false;
+    }
+  }
+  return false;
+}
+
+inline bool function_uses_xla(TF_Function* func, TF_Status* status) {
+  if (function_attr_bool_true(func, "_XlaMustCompile")) {
+    return true;
+  }
+  TF_Buffer* func_def = TF_NewBuffer();
+  TF_FunctionToFunctionDef(func, func_def, status);
+  if (TF_GetCode(status) != TF_OK) {
+    std::string msg = TF_Message(status);
+    TF_DeleteBuffer(func_def);
+    throw deepmd::deepmd_exception("TensorFlow C API Error: " + msg);
+  }
+  const bool uses_xla =
+      func_def->data != NULL &&
+      function_def_uses_xla(static_cast<const char*>(func_def->data),
+                            func_def->length);
+  TF_DeleteBuffer(func_def);
+  return uses_xla;
+}
+
+inline bool graph_uses_xla_compilation(TF_Graph* graph) {
+  size_t pos = 0;
+  while (TF_Operation* oper = TF_GraphNextOperation(graph, &pos)) {
+    const char* op_type = TF_OperationOpType(oper);
+    if ((op_type != NULL && std::strcmp(op_type, "XlaCallModule") == 0) ||
+        operation_attr_bool_true(oper, "_XlaMustCompile")) {
+      return true;
+    }
+  }
+  return false;
+}
+
+inline bool uses_xla_compilation(TF_Graph* graph,
+                                 const std::vector<TF_Function*>& funcs,
+                                 TF_Status* status) {
+  if (graph_uses_xla_compilation(graph)) {
+    return true;
+  }
+  for (TF_Function* func : funcs) {
+    if (function_uses_xla(func, status)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 inline TF_DataType get_data_tensor_type(const std::vector<double>& data) {
   return TF_DOUBLE;
 }
@@ -278,6 +539,7 @@ void deepmd::DeepPotJAX::init(const std::string& model,
   TF_Function** funcs = func_vector.data();
   TF_GraphGetFunctions(graph, funcs, nfuncs, status);
   check_status(status);
+  uses_xla_compilation_ = uses_xla_compilation(graph, func_vector, status);
 
   ctx_opts = TFE_NewContextOptions();
   TFE_ContextOptionsSetConfig(ctx_opts, config.data(), config.size(), status);
@@ -542,16 +804,21 @@ void deepmd::DeepPotJAX::compute(std::vector<ENERGYTYPE>& ener,
   std::vector<double> fparam_double(fparam.begin(), fparam.end());
   std::vector<double> aparam_double(aparam.begin(), aparam.end());
 
-  if (padding_for_nloc != nloc_real) {
-    padding_to_nall = nall_real * PADDING_FACTOR;
-    padding_for_nloc = nloc_real;
-  }
-  while (padding_to_nall < nall_real) {
-    padding_to_nall *= PADDING_FACTOR;
+  int nall_model = nall_real;
+  if (uses_xla_compilation_) {
+    if (padding_for_nloc != nloc_real) {
+      padding_to_nall = nall_real * PADDING_FACTOR;
+      padding_for_nloc = nloc_real;
+    }
+    while (padding_to_nall < nall_real) {
+      padding_to_nall *= PADDING_FACTOR;
+    }
+    nall_model = padding_to_nall;
   }
-  // do padding
-  coord_double.resize(nframes * padding_to_nall * 3, 0.0);
-  atype.resize(nframes * padding_to_nall, -1);
+  // Padding is only useful for XLA-compiled functions; eager TF graphs can use
+  // the exact atom count without shape recompilation churn.
+  coord_double.resize(static_cast<size_t>(nframes) * nall_model * 3, 0.0);
+  atype.resize(static_cast<size_t>(nframes) * nall_model, -1);
 
   TFE_Op* op;
   if (atomic) {
@@ -564,11 +831,11 @@ void deepmd::DeepPotJAX::compute(std::vector<ENERGYTYPE>& ener,
   std::vector<TFE_TensorHandle*> input_list(6);
   std::vector<TF_Tensor*> data_tensor(6);
   // coord
-  std::vector<int64_t> coord_shape = {nframes, padding_to_nall, 3};
+  std::vector<int64_t> coord_shape = {nframes, nall_model, 3};
   input_list[0] =
       add_input(op, coord_double, coord_shape, data_tensor[0], status);
   // atype
-  std::vector<int64_t> atype_shape = {nframes, padding_to_nall};
+  std::vector<int64_t> atype_shape = {nframes, nall_model};
   input_list[1] = add_input(op, atype, atype_shape, data_tensor[1], status);
   // nlist
   if (ago == 0) {
@@ -595,8 +862,8 @@ void deepmd::DeepPotJAX::compute(std::vector<ENERGYTYPE>& ener,
   }
   input_list[2] = add_input(op, nlist, nlist_shape, data_tensor[2], status);
   // mapping; for now, set it to -1, assume it is not used
-  std::vector<int64_t> mapping_shape = {nframes, padding_to_nall};
-  std::vector<int64_t> mapping(nframes * padding_to_nall, -1);
+  std::vector<int64_t> mapping_shape = {nframes, nall_model};
+  std::vector<int64_t> mapping(static_cast<size_t>(nframes) * nall_model, -1);
   // pass mapping if it is given in the neighbor list
   if (lmp_list.mapping) {
     // assume nframes is 1