fix(tabulate): handle empty SE-A neighbors

source/lib/src/gpu/tabulate.cu

@@ -1057,6 +1057,13 @@ void tabulate_fusion_se_a_gpu(FPTYPE* out,
   }
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
+  if (nnei <= 0) {
+    // The descriptor does not carry the empty neighbor dimension, so its
+    // mathematically empty reduction must be materialized explicitly.
+    DPErrcheck(gpuMemset(out, 0, sizeof(FPTYPE) * nloc * MM * last_layer_size));
+    DPErrcheck(gpuDeviceSynchronize());
+    return;
+  }
   tabulate_fusion_se_a_fifth_order_polynomial<FPTYPE, MM, KK>
 #if GOOGLE_CUDA
       <<<nloc, last_layer_size>>>
@@ -1086,7 +1093,7 @@ void tabulate_fusion_se_a_grad_gpu(FPTYPE* dy_dem_x,
                                    const int nnei,
                                    const int last_layer_size,
                                    const bool is_sorted) {
-  if (nloc <= 0) {
+  if (nloc <= 0 || nnei <= 0) {
     return;
   }
   DPErrcheck(gpuGetLastError());
@@ -1128,6 +1135,14 @@ void tabulate_fusion_se_a_grad_grad_gpu(FPTYPE* dz_dy,
   }
   DPErrcheck(gpuGetLastError());
   DPErrcheck(gpuDeviceSynchronize());
+  if (nnei <= 0) {
+    // Unlike the neighbor-shaped inputs, dz_dy remains non-empty and must be
+    // initialized to the zero second derivative of an empty reduction.
+    DPErrcheck(
+        gpuMemset(dz_dy, 0, sizeof(FPTYPE) * nloc * MM * last_layer_size));
+    DPErrcheck(gpuDeviceSynchronize());
+    return;
+  }
   DPErrcheck(gpuMemset(dz_dy, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size));
   tabulate_fusion_se_a_grad_grad_fifth_order_polynomial<FPTYPE, MM, KK>
       <<<nloc, last_layer_size, sizeof(FPTYPE) * MM * last_layer_size>>>(

source/lib/src/tabulate.cc

@@ -165,6 +165,12 @@ void deepmd::tabulate_fusion_se_a_cpu(FPTYPE* out,
                                       const bool is_sorted) {
   bool enable_se_atten = two_embed != nullptr;
   memset(out, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size);
+  // An empty neighbor axis is a valid empty reduction.  Return after
+  // initializing the non-empty descriptor output instead of inspecting the
+  // nonexistent last neighbor below.
+  if (nnei <= 0) {
+    return;
+  }
   const FPTYPE lower = table_info[0];
   const FPTYPE upper = table_info[1];
   const FPTYPE _max = table_info[2];
@@ -246,6 +252,12 @@ void deepmd::tabulate_fusion_se_a_grad_cpu(FPTYPE* dy_dem_x,
                                            const int nnei,
                                            const int last_layer_size,
                                            const bool is_sorted) {
+  // Every gradient output has a zero-sized neighbor axis in this case.  Avoid
+  // both zero-length memory operations on potentially null tensor pointers and
+  // the last-neighbor lookup in the atom loop.
+  if (nnei <= 0) {
+    return;
+  }
   bool enable_se_atten = two_embed != nullptr;
   memset(dy_dem_x, 0, sizeof(FPTYPE) * nloc * nnei);
   memset(dy_dem, 0, sizeof(FPTYPE) * nloc * nnei * 4);
@@ -353,6 +365,11 @@ void deepmd::tabulate_fusion_se_a_grad_grad_cpu(FPTYPE* dz_dy,
                                                 const bool is_sorted) {
   bool enable_se_atten = two_embed != nullptr;
   memset(dz_dy, 0, sizeof(FPTYPE) * nloc * 4 * last_layer_size);
+  // The second-order output retains the descriptor shape, so initialize the
+  // empty reduction to zero before returning.
+  if (nnei <= 0) {
+    return;
+  }
   const FPTYPE lower = table_info[0];
   const FPTYPE upper = table_info[1];
   const FPTYPE _max = table_info[2];

source/lib/tests/test_tabulate_se_a.cc

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: LGPL-3.0-or-later
 #include <gtest/gtest.h>
 
+#include <algorithm>
 #include <cstddef>
 #include <iostream>
 #include <vector>
@@ -757,6 +758,32 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_cpu) {
   }
 }
 
+TEST_F(TestTabulateSeA, empty_neighbors_cpu) {
+  constexpr int empty_nnei = 0;
+  std::vector<double> descriptor(nloc * 4 * last_layer_size, 1.0);
+
+  deepmd::tabulate_fusion_se_a_cpu<double>(
+      descriptor.data(), table.data(), info.data(), nullptr, nullptr, nullptr,
+      nloc, empty_nnei, last_layer_size);
+  for (const double value : descriptor) {
+    EXPECT_DOUBLE_EQ(value, 0.0);
+  }
+
+  // First-order outputs all carry the empty neighbor dimension, so null data
+  // pointers are valid and must not be dereferenced.
+  deepmd::tabulate_fusion_se_a_grad_cpu<double>(
+      nullptr, nullptr, nullptr, table.data(), info.data(), nullptr, nullptr,
+      nullptr, nullptr, nloc, empty_nnei, last_layer_size);
+
+  std::fill(descriptor.begin(), descriptor.end(), 1.0);
+  deepmd::tabulate_fusion_se_a_grad_grad_cpu<double>(
+      descriptor.data(), table.data(), info.data(), nullptr, nullptr, nullptr,
+      nullptr, nullptr, nullptr, nloc, empty_nnei, last_layer_size);
+  for (const double value : descriptor) {
+    EXPECT_DOUBLE_EQ(value, 0.0);
+  }
+}
+
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 TEST_F(TestTabulateSeA, tabulate_fusion_se_a_gpu) {
   std::vector<double> xyz_scatter(nloc * nnei * last_layer_size, 0.0);
@@ -858,6 +885,39 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu) {
   deepmd::delete_device_memory(two_embed_dev);
 }
 
+TEST_F(TestTabulateSeA, empty_neighbors_gpu) {
+  constexpr int empty_nnei = 0;
+  std::vector<double> descriptor(nloc * 4 * last_layer_size, 1.0);
+  double* descriptor_dev = nullptr;
+  deepmd::malloc_device_memory_sync(descriptor_dev, descriptor);
+
+  deepmd::tabulate_fusion_se_a_gpu<double>(descriptor_dev, nullptr, info.data(),
+                                           nullptr, nullptr, nullptr, nloc,
+                                           empty_nnei, last_layer_size);
+  deepmd::memcpy_device_to_host(descriptor_dev, descriptor);
+  for (const double value : descriptor) {
+    EXPECT_DOUBLE_EQ(value, 0.0);
+  }
+
+  // First-order outputs are all empty, and the entry point must return before
+  // launching a kernel or touching their null pointers.
+  deepmd::tabulate_fusion_se_a_grad_gpu<double>(
+      nullptr, nullptr, nullptr, nullptr, info.data(), nullptr, nullptr,
+      nullptr, nullptr, nloc, empty_nnei, last_layer_size);
+
+  std::fill(descriptor.begin(), descriptor.end(), 1.0);
+  deepmd::memcpy_host_to_device(descriptor_dev, descriptor);
+  deepmd::tabulate_fusion_se_a_grad_grad_gpu<double>(
+      descriptor_dev, nullptr, info.data(), nullptr, nullptr, nullptr, nullptr,
+      nullptr, nullptr, nloc, empty_nnei, last_layer_size);
+  deepmd::memcpy_device_to_host(descriptor_dev, descriptor);
+  for (const double value : descriptor) {
+    EXPECT_DOUBLE_EQ(value, 0.0);
+  }
+
+  deepmd::delete_device_memory(descriptor_dev);
+}
+
 TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu_padding_modes) {
   constexpr int test_nloc = 1;
   constexpr int test_nnei = 12;