Skip to content

Commit 58bef11

Browse files
authored
fix(tf): add InvalidArgument compatibility wrapper (#5600)
## Summary - add a TensorFlow-version-gated `deepmd::tf_compat::InvalidArgument` helper - use `absl::InvalidArgumentError` for TensorFlow >= 2.20 and keep `tensorflow::errors::InvalidArgument` for older TensorFlow - route TF custom-op `OP_REQUIRES` InvalidArgument checks through the helper Fixes #5006 @OutisLi Could you review this PR? ## Tests - `uvx ruff==0.15.18 check .` - `uvx ruff==0.15.18 format .` - `DP_ENABLE_PYTORCH=0 uv pip install -e '.[cpu,test]'` with TensorFlow 2.21.0 - `dp --version` - `dp -h`, `dp --tf -h`, `dp --pt -h`, `dp --jax -h`, `dp --pd -h` - `python -c "import deepmd; import deepmd.tf; print('Both interfaces work')"` - `pytest source/tests/tf/test_dp_test.py::TestDPTestEner::test_1frame -v` - standalone C++ build/install with TensorFlow 2.21.0 and PyTorch enabled, using `DEEPMD_BYPASS_TORCH_CUDA_CHECK=ON` for the local CPU PyTorch wheel <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved TensorFlow compatibility across many operators by standardizing invalid-input error handling. * Validation failures now use a framework-compatible error path across supported TensorFlow versions. * Kept the same input checks and messages while making shape/rank mismatch errors more consistent. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent 5082854 commit 58bef11

34 files changed

Lines changed: 1247 additions & 916 deletions

source/op/tf/custom_op.h

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,16 @@
22
#pragma once
33
#include <iostream>
44
#include <string>
5+
#include <utility>
56
#include <vector>
67

8+
#include "tensorflow/core/public/version.h"
9+
10+
#if (TF_MAJOR_VERSION > 2) || (TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 20)
11+
#include "absl/status/status.h"
12+
#include "absl/strings/str_cat.h"
13+
#endif
14+
715
#include "device.h"
816
#include "neighbor_list.h"
917
#include "tensorflow/core/framework/op.h"
@@ -27,6 +35,25 @@ void safe_compute(OpKernelContext* context,
2735
std::function<void(OpKernelContext*)> ff);
2836
};
2937

38+
namespace deepmd {
39+
namespace tf_compat {
40+
#if (TF_MAJOR_VERSION > 2) || (TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 20)
41+
using Status = absl::Status;
42+
#else
43+
using Status = tensorflow::Status;
44+
#endif
45+
46+
template <typename... Args>
47+
inline Status InvalidArgument(Args&&... args) {
48+
#if (TF_MAJOR_VERSION > 2) || (TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 20)
49+
return absl::InvalidArgumentError(absl::StrCat(std::forward<Args>(args)...));
50+
#else
51+
return tensorflow::errors::InvalidArgument(std::forward<Args>(args)...);
52+
#endif
53+
}
54+
} // namespace tf_compat
55+
} // namespace deepmd
56+
3057
template <typename FPTYPE>
3158
void _prepare_coord_nlist_gpu(OpKernelContext* context,
3259
Tensor* tensor_list,

source/op/tf/descrpt.cc

Lines changed: 39 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -67,22 +67,23 @@ class DescrptOp : public OpKernel {
6767

6868
// set size of the sample
6969
OP_REQUIRES(context, (coord_tensor.shape().dims() == 2),
70-
errors::InvalidArgument("Dim of coord should be 2"));
70+
deepmd::tf_compat::InvalidArgument("Dim of coord should be 2"));
7171
OP_REQUIRES(context, (type_tensor.shape().dims() == 2),
72-
errors::InvalidArgument("Dim of type should be 2"));
73-
OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1),
74-
errors::InvalidArgument("Dim of natoms should be 1"));
72+
deepmd::tf_compat::InvalidArgument("Dim of type should be 2"));
73+
OP_REQUIRES(
74+
context, (natoms_tensor.shape().dims() == 1),
75+
deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1"));
7576
OP_REQUIRES(context, (box_tensor.shape().dims() == 2),
76-
errors::InvalidArgument("Dim of box should be 2"));
77+
deepmd::tf_compat::InvalidArgument("Dim of box should be 2"));
7778
OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1),
78-
errors::InvalidArgument("Dim of mesh should be 1"));
79+
deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1"));
7980
OP_REQUIRES(context, (avg_tensor.shape().dims() == 2),
80-
errors::InvalidArgument("Dim of avg should be 2"));
81+
deepmd::tf_compat::InvalidArgument("Dim of avg should be 2"));
8182
OP_REQUIRES(context, (std_tensor.shape().dims() == 2),
82-
errors::InvalidArgument("Dim of std should be 2"));
83+
deepmd::tf_compat::InvalidArgument("Dim of std should be 2"));
8384

8485
OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3),
85-
errors::InvalidArgument(
86+
deepmd::tf_compat::InvalidArgument(
8687
"number of atoms should be larger than (or equal to) 3"));
8788
auto natoms = natoms_tensor.flat<int>();
8889
int nloc = natoms(0);
@@ -91,25 +92,34 @@ class DescrptOp : public OpKernel {
9192
int nsamples = coord_tensor.shape().dim_size(0);
9293

9394
// check the sizes
94-
OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)),
95-
errors::InvalidArgument("number of samples should match"));
96-
OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)),
97-
errors::InvalidArgument("number of samples should match"));
98-
OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)),
99-
errors::InvalidArgument("number of avg should be ntype"));
100-
OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)),
101-
errors::InvalidArgument("number of std should be ntype"));
95+
OP_REQUIRES(
96+
context, (nsamples == type_tensor.shape().dim_size(0)),
97+
deepmd::tf_compat::InvalidArgument("number of samples should match"));
98+
OP_REQUIRES(
99+
context, (nsamples == box_tensor.shape().dim_size(0)),
100+
deepmd::tf_compat::InvalidArgument("number of samples should match"));
101+
OP_REQUIRES(
102+
context, (ntypes == avg_tensor.shape().dim_size(0)),
103+
deepmd::tf_compat::InvalidArgument("number of avg should be ntype"));
104+
OP_REQUIRES(
105+
context, (ntypes == std_tensor.shape().dim_size(0)),
106+
deepmd::tf_compat::InvalidArgument("number of std should be ntype"));
102107

103-
OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)),
104-
errors::InvalidArgument("number of atoms should match"));
105-
OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)),
106-
errors::InvalidArgument("number of atoms should match"));
107-
OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)),
108-
errors::InvalidArgument("number of box should be 9"));
109-
OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)),
110-
errors::InvalidArgument("number of avg should be ndescrpt"));
111-
OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)),
112-
errors::InvalidArgument("number of std should be ndescrpt"));
108+
OP_REQUIRES(
109+
context, (nall * 3 == coord_tensor.shape().dim_size(1)),
110+
deepmd::tf_compat::InvalidArgument("number of atoms should match"));
111+
OP_REQUIRES(
112+
context, (nall == type_tensor.shape().dim_size(1)),
113+
deepmd::tf_compat::InvalidArgument("number of atoms should match"));
114+
OP_REQUIRES(
115+
context, (9 == box_tensor.shape().dim_size(1)),
116+
deepmd::tf_compat::InvalidArgument("number of box should be 9"));
117+
OP_REQUIRES(
118+
context, (ndescrpt == avg_tensor.shape().dim_size(1)),
119+
deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt"));
120+
OP_REQUIRES(
121+
context, (ndescrpt == std_tensor.shape().dim_size(1)),
122+
deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt"));
113123

114124
int nei_mode = 0;
115125
if (mesh_tensor.shape().dim_size(0) == 16) {
@@ -201,10 +211,10 @@ class DescrptOp : public OpKernel {
201211
// }
202212
// int ntypes = max_type_v + 1;
203213
OP_REQUIRES(context, (ntypes == int(sel_a.size())),
204-
errors::InvalidArgument(
214+
deepmd::tf_compat::InvalidArgument(
205215
"number of types should match the length of sel array"));
206216
OP_REQUIRES(context, (ntypes == int(sel_r.size())),
207-
errors::InvalidArgument(
217+
deepmd::tf_compat::InvalidArgument(
208218
"number of types should match the length of sel array"));
209219

210220
for (int kk = 0; kk < nsamples; ++kk) {

source/op/tf/descrpt_se_a_ef.cc

Lines changed: 46 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -69,32 +69,33 @@ class DescrptSeAEfOp : public OpKernel {
6969

7070
// set size of the sample
7171
OP_REQUIRES(context, (coord_tensor.shape().dims() == 2),
72-
errors::InvalidArgument("Dim of coord should be 2"));
72+
deepmd::tf_compat::InvalidArgument("Dim of coord should be 2"));
7373
OP_REQUIRES(context, (type_tensor.shape().dims() == 2),
74-
errors::InvalidArgument("Dim of type should be 2"));
75-
OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1),
76-
errors::InvalidArgument("Dim of natoms should be 1"));
74+
deepmd::tf_compat::InvalidArgument("Dim of type should be 2"));
75+
OP_REQUIRES(
76+
context, (natoms_tensor.shape().dims() == 1),
77+
deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1"));
7778
OP_REQUIRES(context, (box_tensor.shape().dims() == 2),
78-
errors::InvalidArgument("Dim of box should be 2"));
79+
deepmd::tf_compat::InvalidArgument("Dim of box should be 2"));
7980
OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1),
80-
errors::InvalidArgument("Dim of mesh should be 1"));
81+
deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1"));
8182
OP_REQUIRES(context, (ef_tensor.shape().dims() == 2),
82-
errors::InvalidArgument("Dim of ef should be 2"));
83+
deepmd::tf_compat::InvalidArgument("Dim of ef should be 2"));
8384
OP_REQUIRES(context, (avg_tensor.shape().dims() == 2),
84-
errors::InvalidArgument("Dim of avg should be 2"));
85+
deepmd::tf_compat::InvalidArgument("Dim of avg should be 2"));
8586
OP_REQUIRES(context, (std_tensor.shape().dims() == 2),
86-
errors::InvalidArgument("Dim of std should be 2"));
87+
deepmd::tf_compat::InvalidArgument("Dim of std should be 2"));
8788
OP_REQUIRES(
8889
context, (fill_nei_a),
89-
errors::InvalidArgument(
90+
deepmd::tf_compat::InvalidArgument(
9091
"Rotational free descriptor only support the case rcut_a < 0"));
9192
OP_REQUIRES(context, (sec_r.back() == 0),
92-
errors::InvalidArgument(
93+
deepmd::tf_compat::InvalidArgument(
9394
"Rotational free descriptor only support all-angular "
9495
"information: sel_r should be all zero."));
9596

9697
OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3),
97-
errors::InvalidArgument(
98+
deepmd::tf_compat::InvalidArgument(
9899
"number of atoms should be larger than (or equal to) 3"));
99100
auto natoms = natoms_tensor.flat<int>();
100101
int nloc = natoms(0);
@@ -103,29 +104,39 @@ class DescrptSeAEfOp : public OpKernel {
103104
int nsamples = coord_tensor.shape().dim_size(0);
104105

105106
// check the sizes
106-
OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)),
107-
errors::InvalidArgument("number of samples should match"));
108-
OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)),
109-
errors::InvalidArgument("number of samples should match"));
110-
OP_REQUIRES(context, (nsamples == ef_tensor.shape().dim_size(0)),
111-
errors::InvalidArgument("number of samples should match"));
112-
OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)),
113-
errors::InvalidArgument("number of avg should be ntype"));
114-
OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)),
115-
errors::InvalidArgument("number of std should be ntype"));
107+
OP_REQUIRES(
108+
context, (nsamples == type_tensor.shape().dim_size(0)),
109+
deepmd::tf_compat::InvalidArgument("number of samples should match"));
110+
OP_REQUIRES(
111+
context, (nsamples == box_tensor.shape().dim_size(0)),
112+
deepmd::tf_compat::InvalidArgument("number of samples should match"));
113+
OP_REQUIRES(
114+
context, (nsamples == ef_tensor.shape().dim_size(0)),
115+
deepmd::tf_compat::InvalidArgument("number of samples should match"));
116+
OP_REQUIRES(
117+
context, (ntypes == avg_tensor.shape().dim_size(0)),
118+
deepmd::tf_compat::InvalidArgument("number of avg should be ntype"));
119+
OP_REQUIRES(
120+
context, (ntypes == std_tensor.shape().dim_size(0)),
121+
deepmd::tf_compat::InvalidArgument("number of std should be ntype"));
116122

117-
OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)),
118-
errors::InvalidArgument("number of atoms should match"));
119-
OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)),
120-
errors::InvalidArgument("number of atoms should match"));
121-
OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)),
122-
errors::InvalidArgument("number of box should be 9"));
123+
OP_REQUIRES(
124+
context, (nall * 3 == coord_tensor.shape().dim_size(1)),
125+
deepmd::tf_compat::InvalidArgument("number of atoms should match"));
126+
OP_REQUIRES(
127+
context, (nall == type_tensor.shape().dim_size(1)),
128+
deepmd::tf_compat::InvalidArgument("number of atoms should match"));
129+
OP_REQUIRES(
130+
context, (9 == box_tensor.shape().dim_size(1)),
131+
deepmd::tf_compat::InvalidArgument("number of box should be 9"));
123132
OP_REQUIRES(context, (nloc * 3 == ef_tensor.shape().dim_size(1)),
124-
errors::InvalidArgument("number of ef should be 3"));
125-
OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)),
126-
errors::InvalidArgument("number of avg should be ndescrpt"));
127-
OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)),
128-
errors::InvalidArgument("number of std should be ndescrpt"));
133+
deepmd::tf_compat::InvalidArgument("number of ef should be 3"));
134+
OP_REQUIRES(
135+
context, (ndescrpt == avg_tensor.shape().dim_size(1)),
136+
deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt"));
137+
OP_REQUIRES(
138+
context, (ndescrpt == std_tensor.shape().dim_size(1)),
139+
deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt"));
129140

130141
int nei_mode = 0;
131142
if (mesh_tensor.shape().dim_size(0) == 16) {
@@ -208,10 +219,10 @@ class DescrptSeAEfOp : public OpKernel {
208219
// }
209220
// int ntypes = max_type_v + 1;
210221
OP_REQUIRES(context, (ntypes == int(sel_a.size())),
211-
errors::InvalidArgument(
222+
deepmd::tf_compat::InvalidArgument(
212223
"number of types should match the length of sel array"));
213224
OP_REQUIRES(context, (ntypes == int(sel_r.size())),
214-
errors::InvalidArgument(
225+
deepmd::tf_compat::InvalidArgument(
215226
"number of types should match the length of sel array"));
216227

217228
for (int kk = 0; kk < nsamples; ++kk) {

source/op/tf/descrpt_se_a_ef_para.cc

Lines changed: 46 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -69,32 +69,33 @@ class DescrptSeAEfParaOp : public OpKernel {
6969

7070
// set size of the sample
7171
OP_REQUIRES(context, (coord_tensor.shape().dims() == 2),
72-
errors::InvalidArgument("Dim of coord should be 2"));
72+
deepmd::tf_compat::InvalidArgument("Dim of coord should be 2"));
7373
OP_REQUIRES(context, (type_tensor.shape().dims() == 2),
74-
errors::InvalidArgument("Dim of type should be 2"));
75-
OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1),
76-
errors::InvalidArgument("Dim of natoms should be 1"));
74+
deepmd::tf_compat::InvalidArgument("Dim of type should be 2"));
75+
OP_REQUIRES(
76+
context, (natoms_tensor.shape().dims() == 1),
77+
deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1"));
7778
OP_REQUIRES(context, (box_tensor.shape().dims() == 2),
78-
errors::InvalidArgument("Dim of box should be 2"));
79+
deepmd::tf_compat::InvalidArgument("Dim of box should be 2"));
7980
OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1),
80-
errors::InvalidArgument("Dim of mesh should be 1"));
81+
deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1"));
8182
OP_REQUIRES(context, (ef_tensor.shape().dims() == 2),
82-
errors::InvalidArgument("Dim of ef should be 2"));
83+
deepmd::tf_compat::InvalidArgument("Dim of ef should be 2"));
8384
OP_REQUIRES(context, (avg_tensor.shape().dims() == 2),
84-
errors::InvalidArgument("Dim of avg should be 2"));
85+
deepmd::tf_compat::InvalidArgument("Dim of avg should be 2"));
8586
OP_REQUIRES(context, (std_tensor.shape().dims() == 2),
86-
errors::InvalidArgument("Dim of std should be 2"));
87+
deepmd::tf_compat::InvalidArgument("Dim of std should be 2"));
8788
OP_REQUIRES(
8889
context, (fill_nei_a),
89-
errors::InvalidArgument(
90+
deepmd::tf_compat::InvalidArgument(
9091
"Rotational free descriptor only support the case rcut_a < 0"));
9192
OP_REQUIRES(context, (sec_r.back() == 0),
92-
errors::InvalidArgument(
93+
deepmd::tf_compat::InvalidArgument(
9394
"Rotational free descriptor only support all-angular "
9495
"information: sel_r should be all zero."));
9596

9697
OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3),
97-
errors::InvalidArgument(
98+
deepmd::tf_compat::InvalidArgument(
9899
"number of atoms should be larger than (or equal to) 3"));
99100
auto natoms = natoms_tensor.flat<int>();
100101
int nloc = natoms(0);
@@ -103,29 +104,39 @@ class DescrptSeAEfParaOp : public OpKernel {
103104
int nsamples = coord_tensor.shape().dim_size(0);
104105

105106
// check the sizes
106-
OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)),
107-
errors::InvalidArgument("number of samples should match"));
108-
OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)),
109-
errors::InvalidArgument("number of samples should match"));
110-
OP_REQUIRES(context, (nsamples == ef_tensor.shape().dim_size(0)),
111-
errors::InvalidArgument("number of samples should match"));
112-
OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)),
113-
errors::InvalidArgument("number of avg should be ntype"));
114-
OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)),
115-
errors::InvalidArgument("number of std should be ntype"));
107+
OP_REQUIRES(
108+
context, (nsamples == type_tensor.shape().dim_size(0)),
109+
deepmd::tf_compat::InvalidArgument("number of samples should match"));
110+
OP_REQUIRES(
111+
context, (nsamples == box_tensor.shape().dim_size(0)),
112+
deepmd::tf_compat::InvalidArgument("number of samples should match"));
113+
OP_REQUIRES(
114+
context, (nsamples == ef_tensor.shape().dim_size(0)),
115+
deepmd::tf_compat::InvalidArgument("number of samples should match"));
116+
OP_REQUIRES(
117+
context, (ntypes == avg_tensor.shape().dim_size(0)),
118+
deepmd::tf_compat::InvalidArgument("number of avg should be ntype"));
119+
OP_REQUIRES(
120+
context, (ntypes == std_tensor.shape().dim_size(0)),
121+
deepmd::tf_compat::InvalidArgument("number of std should be ntype"));
116122

117-
OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)),
118-
errors::InvalidArgument("number of atoms should match"));
119-
OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)),
120-
errors::InvalidArgument("number of atoms should match"));
121-
OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)),
122-
errors::InvalidArgument("number of box should be 9"));
123+
OP_REQUIRES(
124+
context, (nall * 3 == coord_tensor.shape().dim_size(1)),
125+
deepmd::tf_compat::InvalidArgument("number of atoms should match"));
126+
OP_REQUIRES(
127+
context, (nall == type_tensor.shape().dim_size(1)),
128+
deepmd::tf_compat::InvalidArgument("number of atoms should match"));
129+
OP_REQUIRES(
130+
context, (9 == box_tensor.shape().dim_size(1)),
131+
deepmd::tf_compat::InvalidArgument("number of box should be 9"));
123132
OP_REQUIRES(context, (nloc * 3 == ef_tensor.shape().dim_size(1)),
124-
errors::InvalidArgument("number of ef should be 3"));
125-
OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)),
126-
errors::InvalidArgument("number of avg should be ndescrpt"));
127-
OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)),
128-
errors::InvalidArgument("number of std should be ndescrpt"));
133+
deepmd::tf_compat::InvalidArgument("number of ef should be 3"));
134+
OP_REQUIRES(
135+
context, (ndescrpt == avg_tensor.shape().dim_size(1)),
136+
deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt"));
137+
OP_REQUIRES(
138+
context, (ndescrpt == std_tensor.shape().dim_size(1)),
139+
deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt"));
129140

130141
int nei_mode = 0;
131142
if (mesh_tensor.shape().dim_size(0) == 16) {
@@ -208,10 +219,10 @@ class DescrptSeAEfParaOp : public OpKernel {
208219
// }
209220
// int ntypes = max_type_v + 1;
210221
OP_REQUIRES(context, (ntypes == int(sel_a.size())),
211-
errors::InvalidArgument(
222+
deepmd::tf_compat::InvalidArgument(
212223
"number of types should match the length of sel array"));
213224
OP_REQUIRES(context, (ntypes == int(sel_r.size())),
214-
errors::InvalidArgument(
225+
deepmd::tf_compat::InvalidArgument(
215226
"number of types should match the length of sel array"));
216227

217228
for (int kk = 0; kk < nsamples; ++kk) {

0 commit comments

Comments
 (0)