[BUG] Inference with Python interface DPLR model raises error.

[Kehan-Cai-nanako]

Bug summary

This bug is only specific to the python interface and the compressed Deep Dipole model.

We trained a Deep Potential long range model based on a compressed Deep Dipole model, and used this model to do inference.

If we put it in lammps, it runs fine.
If we use the python interface to do inference, it will raise error.

If we replace this model with one that is based on an uncompressed Deep Dipole model, both interfaces work correctly.

DeePMD-kit Version

DeePMD-kit v2.2.3.dev18+g9e2dfdf1

TensorFlow Version

2.9.0

How did you download the software?

Built from source

Input Files, Running Commands, Error Log, etc.

dplr.zip

Error Log:

Traceback (most recent call last):
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/tensorflow/python/ops/gradients_util.py", line 611, in _GradientsHelper
grad_fn = ops.get_gradient_function(op)
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/tensorflow/python/framework/ops.py", line 2838, in get_gradient_function
return gradient_registry.lookup(op_type)
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/tensorflow/python/framework/registry.py", line 95, in lookup
raise LookupError(
LookupError: gradient registry has no entry for: TabulateFusionSeA

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/scratch/gpfs/kehanc/StudyFold/ResearchFold/Relaxor_Ferroelectrics/preliminary/dplr/prof_1/dpt_infer.py", line 28, in
tensor = dp.eval(coord, cell, types)
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/deepmd/infer/deep_pot.py", line 387, in eval
me, mf, mv = self.dm.eval(coords, cells, atom_types)
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/deepmd/infer/data_modifier.py", line 287, in eval
self.force, self.virial, self.av = self.build_fv_graph()
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/deepmd/infer/data_modifier.py", line 111, in build_fv_graph
return self._build_fv_graph_inner()
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/deepmd/infer/data_modifier.py", line 153, in _build_fv_graph_inner
[self.t_ef_d] = tf.gradients(
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/tensorflow/python/ops/gradients_impl.py", line 165, in gradients
return gradients_util._GradientsHelper(
File "/tigress/yifanl/usr/licensed/anaconda3/2021.11/envs/dpdev1/lib/python3.10/site-packages/tensorflow/python/ops/gradients_util.py", line 639, in _GradientsHelper
raise LookupError(
LookupError: No gradient defined for operation'dipole_charge/filter_type_0/TabulateFusionSeA_3' (op type: TabulateFusionSeA). In general every operation must have an associated @tf.RegisterGradient for correct autodiff, which this op is lacking. If you want to pretend this operation is a constant in your program, you may insert tf.stop_gradient. This can be useful to silence the error in cases where you know gradients are not needed, e.g. the forward pass of tf.custom_gradient. Please see more details in https://www.tensorflow.org/api_docs/python/tf/custom_gradient.

Steps to Reproduce

Run LAMMPS:
lmp -v TEMP 50 -in in_npt.lammps

Run python script:
python dpt_infer.py

Further Information, Files, and Links

No response