Op type not registered 'XlaSharding' in binary running

[Cynthia-0807]

Summary

When I run the following comand,

cd ~/software/DMFF/backend
python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100 --pdb ../examples/water_fullpol/water_dimer.pdb --model ./openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/water_dimer.pdb --has_aux True

it reported the error:

2024-03-18 22:02:19.504317: I tensorflow/cc/saved_model/reader.cc:31] Reading SavedModel from: ./openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/output
2024-03-18 22:02:19.505456: I tensorflow/cc/saved_model/reader.cc:54] Reading meta graph with tags { serve }
2024-03-18 22:02:19.506154: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.1 SSE4.2 AVX AVX2 FMA
2024-03-18 22:02:19.509908: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 2112005000 Hz
2024-03-18 22:02:19.511792: I tensorflow/compiler/xla/service/service.cc:150] XLA service 0x56436b07f550 executing computations on platform Host. Devices:
2024-03-18 22:02:19.511817: I tensorflow/compiler/xla/service/service.cc:158]   StreamExecutor device (0): <undefined>, <undefined>
2024-03-18 22:02:19.519027: I tensorflow/cc/saved_model/loader.cc:182] Restoring SavedModel bundle.
2024-03-18 22:02:19.524285: I tensorflow/cc/saved_model/loader.cc:132] Running initialization op on SavedModel bundle.
2024-03-18 22:02:19.526928: I tensorflow/cc/saved_model/loader.cc:285] SavedModel load for tags { serve }; Status: success. Took 22613 microseconds.
Running dynamics
2024-03-18 22:02:19.541867: W tensorflow/core/kernels/partitioned_function_ops.cc:197] Grappler optimization failed. Error: Op type not registered 'XlaSharding' in binary running on DESKTOP-87V2MP7. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
2024-03-18 22:02:19.542614: W tensorflow/core/framework/op_kernel.cc:1401] OP_REQUIRES failed at partitioned_function_ops.cc:118 : Not found: Op type not registered 'XlaSharding' in binary running on DESKTOP-87V2MP7. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
Traceback (most recent call last):
  File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/OpenMMDMFFPlugin/tests/test_dmff_plugin_nve.py", line 109, in <module>
    test_dmff_nve(nsteps=nsteps, time_step=time_step, pdb_file=pdb, model_dir=model_dir, platform_name=platform_name, has_aux=args.has_aux)
  File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/OpenMMDMFFPlugin/tests/test_dmff_plugin_nve.py", line 64, in test_dmff_nve
    sim.step(nsteps)
  File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/openmm/app/simulation.py", line 141, in step
    self._simulate(endStep=self.currentStep+steps)
  File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/openmm/app/simulation.py", line 210, in _simulate
    self.integrator.step(stepsToGo)
  File "/home/wyq/software/anaconda3/envs/dmff_omm/lib/python3.9/site-packages/openmm/openmm.py", line 3092, in step
    return _openmm.VerletIntegrator_step(self, steps)
openmm.OpenMMException: Op type not registered 'XlaSharding' in binary running on DESKTOP-87V2MP7. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
         [[{{node StatefulPartitionedCall}}]]
         [[{{node StatefulPartitionedCall}}]]

Could someone please take a look for me?

DMFF Version

1.0.0

JAX Version

0.4.25

OpenMM Version

7.7.0

Python Version, CUDA Version, GCC Version, Operating System Version etc

No response

Details

the version of tensorflow is 2.16.1

[Cynthia-0807]

The problem may be from unmatching the version of libtensorflow_cc to one of tensorflow.When I changed the verison of libtensorflow_cc,it worked.