diff --git a/source/op/tf/custom_op.h b/source/op/tf/custom_op.h index baf95f3fa3..b133c0e46b 100644 --- a/source/op/tf/custom_op.h +++ b/source/op/tf/custom_op.h @@ -2,8 +2,16 @@ #pragma once #include #include +#include #include +#include "tensorflow/core/public/version.h" + +#if (TF_MAJOR_VERSION > 2) || (TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 20) +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" +#endif + #include "device.h" #include "neighbor_list.h" #include "tensorflow/core/framework/op.h" @@ -27,6 +35,25 @@ void safe_compute(OpKernelContext* context, std::function ff); }; +namespace deepmd { +namespace tf_compat { +#if (TF_MAJOR_VERSION > 2) || (TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 20) +using Status = absl::Status; +#else +using Status = tensorflow::Status; +#endif + +template +inline Status InvalidArgument(Args&&... args) { +#if (TF_MAJOR_VERSION > 2) || (TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 20) + return absl::InvalidArgumentError(absl::StrCat(std::forward(args)...)); +#else + return tensorflow::errors::InvalidArgument(std::forward(args)...); +#endif +} +} // namespace tf_compat +} // namespace deepmd + template void _prepare_coord_nlist_gpu(OpKernelContext* context, Tensor* tensor_list, diff --git a/source/op/tf/descrpt.cc b/source/op/tf/descrpt.cc index db3b0ca8e5..37d0764e47 100644 --- a/source/op/tf/descrpt.cc +++ b/source/op/tf/descrpt.cc @@ -67,22 +67,23 @@ class DescrptOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); int nloc = natoms(0); @@ -91,25 +92,34 @@ class DescrptOp : public OpKernel { int nsamples = coord_tensor.shape().dim_size(0); // check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); int nei_mode = 0; if (mesh_tensor.shape().dim_size(0) == 16) { @@ -201,10 +211,10 @@ class DescrptOp : public OpKernel { // } // int ntypes = max_type_v + 1; OP_REQUIRES(context, (ntypes == int(sel_a.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); OP_REQUIRES(context, (ntypes == int(sel_r.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); for (int kk = 0; kk < nsamples; ++kk) { diff --git a/source/op/tf/descrpt_se_a_ef.cc b/source/op/tf/descrpt_se_a_ef.cc index 18dda3d8b0..170c3bb5e8 100644 --- a/source/op/tf/descrpt_se_a_ef.cc +++ b/source/op/tf/descrpt_se_a_ef.cc @@ -69,32 +69,33 @@ class DescrptSeAEfOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (ef_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of ef should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of ef should be 2")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES( context, (fill_nei_a), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support the case rcut_a < 0")); OP_REQUIRES(context, (sec_r.back() == 0), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support all-angular " "information: sel_r should be all zero.")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); int nloc = natoms(0); @@ -103,29 +104,39 @@ class DescrptSeAEfOp : public OpKernel { int nsamples = coord_tensor.shape().dim_size(0); // check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == ef_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == ef_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); OP_REQUIRES(context, (nloc * 3 == ef_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of ef should be 3")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + deepmd::tf_compat::InvalidArgument("number of ef should be 3")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); int nei_mode = 0; if (mesh_tensor.shape().dim_size(0) == 16) { @@ -208,10 +219,10 @@ class DescrptSeAEfOp : public OpKernel { // } // int ntypes = max_type_v + 1; OP_REQUIRES(context, (ntypes == int(sel_a.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); OP_REQUIRES(context, (ntypes == int(sel_r.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); for (int kk = 0; kk < nsamples; ++kk) { diff --git a/source/op/tf/descrpt_se_a_ef_para.cc b/source/op/tf/descrpt_se_a_ef_para.cc index 0f34de3f4f..d7c4dc24b1 100644 --- a/source/op/tf/descrpt_se_a_ef_para.cc +++ b/source/op/tf/descrpt_se_a_ef_para.cc @@ -69,32 +69,33 @@ class DescrptSeAEfParaOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (ef_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of ef should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of ef should be 2")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES( context, (fill_nei_a), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support the case rcut_a < 0")); OP_REQUIRES(context, (sec_r.back() == 0), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support all-angular " "information: sel_r should be all zero.")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); int nloc = natoms(0); @@ -103,29 +104,39 @@ class DescrptSeAEfParaOp : public OpKernel { int nsamples = coord_tensor.shape().dim_size(0); // check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == ef_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == ef_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); OP_REQUIRES(context, (nloc * 3 == ef_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of ef should be 3")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + deepmd::tf_compat::InvalidArgument("number of ef should be 3")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); int nei_mode = 0; if (mesh_tensor.shape().dim_size(0) == 16) { @@ -208,10 +219,10 @@ class DescrptSeAEfParaOp : public OpKernel { // } // int ntypes = max_type_v + 1; OP_REQUIRES(context, (ntypes == int(sel_a.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); OP_REQUIRES(context, (ntypes == int(sel_r.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); for (int kk = 0; kk < nsamples; ++kk) { diff --git a/source/op/tf/descrpt_se_a_ef_vert.cc b/source/op/tf/descrpt_se_a_ef_vert.cc index b4eb30d9ee..a026495b3b 100644 --- a/source/op/tf/descrpt_se_a_ef_vert.cc +++ b/source/op/tf/descrpt_se_a_ef_vert.cc @@ -69,32 +69,33 @@ class DescrptSeAEfVertOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (ef_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of ef should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of ef should be 2")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES( context, (fill_nei_a), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support the case rcut_a < 0")); OP_REQUIRES(context, (sec_r.back() == 0), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support all-angular " "information: sel_r should be all zero.")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); int nloc = natoms(0); @@ -103,29 +104,39 @@ class DescrptSeAEfVertOp : public OpKernel { int nsamples = coord_tensor.shape().dim_size(0); // check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == ef_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == ef_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); OP_REQUIRES(context, (nloc * 3 == ef_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of ef should be 3")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + deepmd::tf_compat::InvalidArgument("number of ef should be 3")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); int nei_mode = 0; if (mesh_tensor.shape().dim_size(0) == 16) { @@ -208,10 +219,10 @@ class DescrptSeAEfVertOp : public OpKernel { // } // int ntypes = max_type_v + 1; OP_REQUIRES(context, (ntypes == int(sel_a.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); OP_REQUIRES(context, (ntypes == int(sel_r.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); for (int kk = 0; kk < nsamples; ++kk) { diff --git a/source/op/tf/descrpt_se_a_mask.cc b/source/op/tf/descrpt_se_a_mask.cc index 7f8bcd9411..f04237d067 100644 --- a/source/op/tf/descrpt_se_a_mask.cc +++ b/source/op/tf/descrpt_se_a_mask.cc @@ -63,20 +63,23 @@ class DescrptSeAMaskOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); + OP_REQUIRES(context, (type_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument( + "Dim of type for se_e2_a_mask op should be 2")); OP_REQUIRES( - context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type for se_e2_a_mask op should be 2")); - OP_REQUIRES(context, (mask_matrix_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of mask matrix should be 2")); + context, (mask_matrix_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of mask matrix should be 2")); int nsamples = coord_tensor.shape().dim_size(0); // check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == mask_matrix_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == mask_matrix_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); // Set n_descrpt for each atom. Include 1/rr, cos(theta), cos(phi), sin(phi) int n_descrpt = 4; @@ -85,12 +88,12 @@ class DescrptSeAMaskOp : public OpKernel { auto natoms = natoms_tensor.flat(); total_atom_num = natoms(1); // check the sizes - OP_REQUIRES(context, - (total_atom_num * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, - (total_atom_num == mask_matrix_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (total_atom_num * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (total_atom_num == mask_matrix_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); // Create an output tensor TensorShape descrpt_shape; diff --git a/source/op/tf/ewald_recp.cc b/source/op/tf/ewald_recp.cc index dcad204467..8bfefc67d7 100644 --- a/source/op/tf/ewald_recp.cc +++ b/source/op/tf/ewald_recp.cc @@ -43,31 +43,33 @@ class EwaldRecpOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (coord_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of coord should be 1")); - OP_REQUIRES(context, (charge_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of charge should be 1")); - OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) == 1), - errors::InvalidArgument("size of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 1")); + OP_REQUIRES( + context, (charge_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of charge should be 1")); + OP_REQUIRES( + context, (natoms_tensor.shape().dim_size(0) == 1), + deepmd::tf_compat::InvalidArgument("size of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of box should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 1")); auto natoms = natoms_tensor.flat(); int nloc = natoms(0); int nsamples = coord_tensor.shape().dim_size(0) / (nloc * 3); // check the sizes - OP_REQUIRES( - context, - (static_cast(nsamples) * nloc * 3 == - coord_tensor.shape().dim_size(0)), - errors::InvalidArgument("coord number of samples should match")); - OP_REQUIRES( - context, - (static_cast(nsamples) * nloc * 1 == - charge_tensor.shape().dim_size(0)), - errors::InvalidArgument("charge number of samples should match")); - OP_REQUIRES( - context, (nsamples * 9 == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("box number of samples should match")); + OP_REQUIRES(context, + (static_cast(nsamples) * nloc * 3 == + coord_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument( + "coord number of samples should match")); + OP_REQUIRES(context, + (static_cast(nsamples) * nloc * 1 == + charge_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument( + "charge number of samples should match")); + OP_REQUIRES(context, (nsamples * 9 == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument( + "box number of samples should match")); // Create an output tensor TensorShape energy_shape; diff --git a/source/op/tf/map_aparam.cc b/source/op/tf/map_aparam.cc index 7ac3b48a4f..e5d8059146 100644 --- a/source/op/tf/map_aparam.cc +++ b/source/op/tf/map_aparam.cc @@ -34,15 +34,17 @@ class MapAparamOp : public OpKernel { const Tensor& natoms_tensor = context->input(context_input_index++); // set size of the sample - OP_REQUIRES(context, (aparam_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of aparam should be 2")); + OP_REQUIRES( + context, (aparam_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of aparam should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -53,10 +55,12 @@ class MapAparamOp : public OpKernel { int numb_aparam = nall > 0 ? aparam_tensor.shape().dim_size(1) / nall : 0; // check the sizes - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape output_shape; diff --git a/source/op/tf/neighbor_stat.cc b/source/op/tf/neighbor_stat.cc index 26f13b0c84..fdb1df2456 100644 --- a/source/op/tf/neighbor_stat.cc +++ b/source/op/tf/neighbor_stat.cc @@ -46,33 +46,39 @@ class NeighborStatOp : public OpKernel { const Tensor& mesh_tensor = context->input(context_input_index++); OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); int nloc = natoms_tensor.flat().data()[0]; int nall = natoms_tensor.flat().data()[1]; int nsamples = coord_tensor.shape().dim_size(0); int ntypes = natoms_tensor.shape().dim_size(0) - 2; // check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); DeviceFunctor()(device, context->eigen_device()); int nei_mode = 0; if (mesh_tensor.shape().dim_size(0) == 6 || diff --git a/source/op/tf/pair_tab.cc b/source/op/tf/pair_tab.cc index 0b04390d5f..e14b5c1e80 100644 --- a/source/op/tf/pair_tab.cc +++ b/source/op/tf/pair_tab.cc @@ -52,23 +52,26 @@ class PairTabOp : public OpKernel { const Tensor& scale_tensor = context->input(tmp_idx++); // set size of the sample - OP_REQUIRES(context, (table_info_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of table_info should be 1")); - OP_REQUIRES(context, (table_data_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of table_data should be 1")); + OP_REQUIRES( + context, (table_info_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of table_info should be 1")); + OP_REQUIRES( + context, (table_data_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of table_data should be 1")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); OP_REQUIRES(context, (rij_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (scale_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of scale should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of scale should be 2")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -80,25 +83,31 @@ class PairTabOp : public OpKernel { assert(sel_r.size() == ntypes); // check the sizes - OP_REQUIRES(context, (nframes == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == rij_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("shape of type should be nall")); OP_REQUIRES( - context, - (3 * static_cast(nnei) * nloc == - rij_tensor.shape().dim_size(1)), - errors::InvalidArgument("shape of rij should be 3 * nloc * nnei")); + context, (nframes == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == rij_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("shape of type should be nall")); + OP_REQUIRES(context, + (3 * static_cast(nnei) * nloc == + rij_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "shape of rij should be 3 * nloc * nnei")); OP_REQUIRES( context, (static_cast(nnei) * nloc == nlist_tensor.shape().dim_size(1)), - errors::InvalidArgument("shape of nlist should be nloc * nnei")); - OP_REQUIRES(context, (nloc == scale_tensor.shape().dim_size(1)), - errors::InvalidArgument("shape of scale should be nloc")); + deepmd::tf_compat::InvalidArgument( + "shape of nlist should be nloc * nnei")); + OP_REQUIRES( + context, (nloc == scale_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("shape of scale should be nloc")); // Create an output tensor TensorShape energy_shape; @@ -133,7 +142,7 @@ class PairTabOp : public OpKernel { auto virial = virial_tensor->matrix(); OP_REQUIRES(context, (ntypes == int(table_info(3) + 0.1)), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "ntypes provided in table does not match deeppot")); int nspline = table_info(2) + 0.1; int tab_stride = 4 * nspline; diff --git a/source/op/tf/pairwise.cc b/source/op/tf/pairwise.cc index ba1e5e6475..eff7b8ffe2 100644 --- a/source/op/tf/pairwise.cc +++ b/source/op/tf/pairwise.cc @@ -45,9 +45,10 @@ class PairwiseIdxOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (idxs_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of idxs should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of idxs should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); auto idxs = idxs_tensor.matrix(); int nframes = idxs_tensor.shape().dim_size(0); @@ -55,7 +56,7 @@ class PairwiseIdxOp : public OpKernel { int nloc = natoms(0); int nall = natoms(1); OP_REQUIRES(context, nframes > 0, - errors::InvalidArgument("nframes should be > 0")); + deepmd::tf_compat::InvalidArgument("nframes should be > 0")); std::vector> forward_qm_maps, backward_qm_maps, forward_qmmm_maps, backward_qmmm_maps; @@ -237,9 +238,10 @@ class ConvertForwardMapOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (sub_forward_map_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of idxs should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of idxs should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); auto sub_forward_map = sub_forward_map_tensor.matrix(); int sub_nframes = sub_forward_map_tensor.shape().dim_size(0); diff --git a/source/op/tf/prod_env_mat_multi_device.cc b/source/op/tf/prod_env_mat_multi_device.cc index e374102224..e57c08356e 100644 --- a/source/op/tf/prod_env_mat_multi_device.cc +++ b/source/op/tf/prod_env_mat_multi_device.cc @@ -364,25 +364,26 @@ class ProdEnvMatAOp : public OpKernel { // set size of the sample. assume 't' is [[[1, 1, 1], [2, 2, 2]], [[3, 3, // 3], [4, 4, 4]]], then shape(t) ==> [2, 2, 3] OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES(context, (sec_r.back() == 0), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support all-angular " "information: sel_r should be all zero.")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); DeviceFunctor()(device, context->eigen_device()); const int* natoms = natoms_tensor.flat().data(); @@ -392,31 +393,40 @@ class ProdEnvMatAOp : public OpKernel { natoms_tensor.shape().dim_size(0) - 2; // nloc and nall mean something. int nsamples = coord_tensor.shape().dim_size(0); //// check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); - - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); + + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); OP_REQUIRES(context, (ntypes == int(sel_a.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); OP_REQUIRES(context, (ntypes == int(sel_r.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); int nei_mode = 0; @@ -680,21 +690,22 @@ class ProdEnvMatROp : public OpKernel { const Tensor& std_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); DeviceFunctor()(device, context->eigen_device()); const int* natoms = natoms_tensor.flat().data(); @@ -706,24 +717,33 @@ class ProdEnvMatROp : public OpKernel { // //// check the sizes // check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); int nei_mode = 0; bool b_nlist_map = false; @@ -995,25 +1015,26 @@ class ProdEnvMatAMixOp : public OpKernel { // set size of the sample. assume 't' is [[[1, 1, 1], [2, 2, 2]], [[3, 3, // 3], [4, 4, 4]]], then shape(t) ==> [2, 2, 3] OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES(context, (sec_r.back() == 0), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support all-angular " "information: sel_r should be all zero.")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); DeviceFunctor()(device, context->eigen_device()); const int* natoms = natoms_tensor.flat().data(); @@ -1022,31 +1043,40 @@ class ProdEnvMatAMixOp : public OpKernel { int ntypes = natoms_tensor.shape().dim_size(0) - 2; int nsamples = coord_tensor.shape().dim_size(0); //// check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); - - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); + + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); OP_REQUIRES(context, (1 == int(sel_a.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "the length of sel array should be 1 in this op")); OP_REQUIRES(context, (1 == int(sel_r.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "the length of sel array should be 1 in this op")); int nei_mode = 0; @@ -1716,7 +1746,7 @@ void _prepare_coord_nlist_gpu(OpKernelContext* context, deepmd::env_mat_nbor_update(inlist_temp, inlist, max_nbor_size, nbor_list_dev, fake_mesh_dev, 16); OP_REQUIRES(context, (max_numneigh(inlist_temp) <= max_nbor_size), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Assert failed, max neighbor size of atom(lammps) " + std::to_string(max_numneigh(inlist_temp)) + " is larger than " + std::to_string(max_nbor_size) + @@ -1730,7 +1760,7 @@ void _prepare_coord_nlist_gpu(OpKernelContext* context, nbor_list_dev, mesh_tensor_data, mesh_tensor_size); OP_REQUIRES(context, (max_numneigh(inlist_temp) <= max_nbor_size), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Assert failed, max neighbor size of atom(lammps) " + std::to_string(max_numneigh(inlist_temp)) + " is larger than " + std::to_string(max_nbor_size) + diff --git a/source/op/tf/prod_env_mat_multi_device_nvnmd.cc b/source/op/tf/prod_env_mat_multi_device_nvnmd.cc index 57390077ef..95bcb71b3c 100644 --- a/source/op/tf/prod_env_mat_multi_device_nvnmd.cc +++ b/source/op/tf/prod_env_mat_multi_device_nvnmd.cc @@ -342,25 +342,26 @@ class ProdEnvMatANvnmdQuantizeOp : public OpKernel { // set size of the sample. assume 't' is [[[1, 1, 1], [2, 2, 2]], [[3, 3, // 3], [4, 4, 4]]], then shape(t) ==> [2, 2, 3] OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES(context, (sec_r.back() == 0), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support all-angular " "information: sel_r should be all zero.")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); DeviceFunctor()(device, context->eigen_device()); const int* natoms = natoms_tensor.flat().data(); @@ -370,31 +371,40 @@ class ProdEnvMatANvnmdQuantizeOp : public OpKernel { natoms_tensor.shape().dim_size(0) - 2; // nloc and nall mean something. int nsamples = coord_tensor.shape().dim_size(0); //// check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); - - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); + + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); OP_REQUIRES(context, (ntypes == int(sel_a.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); OP_REQUIRES(context, (ntypes == int(sel_r.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of types should match the length of sel array")); int nei_mode = 0; @@ -585,25 +595,26 @@ class ProdEnvMatAMixNvnmdQuantizeOp : public OpKernel { // set size of the sample. assume 't' is [[[1, 1, 1], [2, 2, 2]], [[3, 3, // 3], [4, 4, 4]]], then shape(t) ==> [2, 2, 3] OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of coord should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of coord should be 2")); OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (box_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of box should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of box should be 2")); OP_REQUIRES(context, (mesh_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of mesh should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of mesh should be 1")); OP_REQUIRES(context, (avg_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of avg should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of avg should be 2")); OP_REQUIRES(context, (std_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of std should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of std should be 2")); OP_REQUIRES(context, (sec_r.back() == 0), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "Rotational free descriptor only support all-angular " "information: sel_r should be all zero.")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); DeviceFunctor()(device, context->eigen_device()); const int* natoms = natoms_tensor.flat().data(); @@ -612,31 +623,40 @@ class ProdEnvMatAMixNvnmdQuantizeOp : public OpKernel { int ntypes = natoms_tensor.shape().dim_size(0) - 2; int nsamples = coord_tensor.shape().dim_size(0); //// check the sizes - OP_REQUIRES(context, (nsamples == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nsamples == box_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (ntypes == avg_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of avg should be ntype")); - OP_REQUIRES(context, (ntypes == std_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of std should be ntype")); - - OP_REQUIRES(context, (nall * 3 == coord_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of atoms should match")); - OP_REQUIRES(context, (9 == box_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of box should be 9")); - OP_REQUIRES(context, (ndescrpt == avg_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of avg should be ndescrpt")); - OP_REQUIRES(context, (ndescrpt == std_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of std should be ndescrpt")); + OP_REQUIRES( + context, (nsamples == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nsamples == box_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (ntypes == avg_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of avg should be ntype")); + OP_REQUIRES( + context, (ntypes == std_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of std should be ntype")); + + OP_REQUIRES( + context, (nall * 3 == coord_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of atoms should match")); + OP_REQUIRES( + context, (9 == box_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of box should be 9")); + OP_REQUIRES( + context, (ndescrpt == avg_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of avg should be ndescrpt")); + OP_REQUIRES( + context, (ndescrpt == std_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of std should be ndescrpt")); OP_REQUIRES(context, (1 == int(sel_a.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "the length of sel array should be 1 in this op")); OP_REQUIRES(context, (1 == int(sel_r.size())), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "the length of sel array should be 1 in this op")); int nei_mode = 0; diff --git a/source/op/tf/prod_force.cc b/source/op/tf/prod_force.cc index 20269ebef3..9f0d1af8de 100644 --- a/source/op/tf/prod_force.cc +++ b/source/op/tf/prod_force.cc @@ -39,19 +39,22 @@ class ProdForceOp : public OpKernel { const Tensor& natoms_tensor = context->input(4); // set size of the sample - OP_REQUIRES(context, (net_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + OP_REQUIRES( + context, (net_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); OP_REQUIRES(context, (axis_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of axis should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of axis should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -62,22 +65,27 @@ class ProdForceOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == in_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == axis_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == in_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == axis_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 12 == in_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); OP_REQUIRES( - context, (nloc * 4 == axis_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of axis type+id should match 2+2")); + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); + OP_REQUIRES(context, (nloc * 4 == axis_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "number of axis type+id should match 2+2")); // Create an output tensor TensorShape force_shape; diff --git a/source/op/tf/prod_force_grad.cc b/source/op/tf/prod_force_grad.cc index acfe9145fe..f5b9295643 100644 --- a/source/op/tf/prod_force_grad.cc +++ b/source/op/tf/prod_force_grad.cc @@ -46,20 +46,23 @@ class ProdForceGradOp : public OpKernel { TensorShape axis_shape = axis_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); OP_REQUIRES(context, (axis_shape.dims() == 2), - errors::InvalidArgument("Dim of axis should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of axis should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -69,27 +72,33 @@ class ProdForceGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == axis_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (nloc * 3 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == axis_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 12 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); OP_REQUIRES( - context, (nloc * 4 == axis_shape.dim_size(1)), - errors::InvalidArgument("number of axis type+id should be 2+2")); + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); + OP_REQUIRES(context, (nloc * 4 == axis_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "number of axis type+id should be 2+2")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_force_grad_multi_device.cc b/source/op/tf/prod_force_grad_multi_device.cc index ee8a29732d..6678c1b628 100644 --- a/source/op/tf/prod_force_grad_multi_device.cc +++ b/source/op/tf/prod_force_grad_multi_device.cc @@ -53,18 +53,21 @@ class ProdForceSeAGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -74,21 +77,26 @@ class ProdForceSeAGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (nloc * 3 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape grad_net_shape; @@ -166,18 +174,21 @@ class ProdForceSeRGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -187,19 +198,23 @@ class ProdForceSeRGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (nloc * 3 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_force_multi_device.cc b/source/op/tf/prod_force_multi_device.cc index d48749faa5..ae66f13503 100644 --- a/source/op/tf/prod_force_multi_device.cc +++ b/source/op/tf/prod_force_multi_device.cc @@ -74,16 +74,19 @@ class ProdForceSeAOp : public OpKernel { const Tensor& nlist_tensor = context->input(context_input_index++); const Tensor& natoms_tensor = context->input(context_input_index++); // set size of the sample - OP_REQUIRES(context, (net_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + OP_REQUIRES( + context, (net_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); const int* natoms = natoms_tensor.flat().data(); int nloc = natoms[0]; @@ -92,14 +95,17 @@ class ProdForceSeAOp : public OpKernel { int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == in_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == in_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); OP_REQUIRES( context, (int_64(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); // Create an output tensor TensorShape force_shape; force_shape.AddDim(nframes); @@ -175,16 +181,19 @@ class ProdForceSeROp : public OpKernel { const Tensor& nlist_tensor = context->input(context_input_index++); const Tensor& natoms_tensor = context->input(context_input_index++); // set size of the sample - OP_REQUIRES(context, (net_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + OP_REQUIRES( + context, (net_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); const int* natoms = natoms_tensor.flat().data(); int nloc = natoms[0]; @@ -193,14 +202,17 @@ class ProdForceSeROp : public OpKernel { int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == in_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == in_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); // Create an output tensor TensorShape force_shape; force_shape.AddDim(nframes); diff --git a/source/op/tf/prod_force_se_a_grad.cc b/source/op/tf/prod_force_se_a_grad.cc index 05e26b5058..5a270f88f4 100644 --- a/source/op/tf/prod_force_se_a_grad.cc +++ b/source/op/tf/prod_force_se_a_grad.cc @@ -46,18 +46,21 @@ class ProdForceSeAGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -67,22 +70,27 @@ class ProdForceSeAGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (nloc * 3 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_force_se_a_mask.cc b/source/op/tf/prod_force_se_a_mask.cc index 6c938f88e0..ae28aa670e 100644 --- a/source/op/tf/prod_force_se_a_mask.cc +++ b/source/op/tf/prod_force_se_a_mask.cc @@ -36,14 +36,17 @@ class ProdForceSeAMaskOp : public OpKernel { const Tensor& nlist_tensor = context->input(3); // set size of the sample - OP_REQUIRES(context, (net_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); - OP_REQUIRES(context, (mask_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of mask matrix should be 2")); + OP_REQUIRES( + context, (net_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); + OP_REQUIRES( + context, (mask_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of mask matrix should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = total_atom_num; @@ -52,14 +55,17 @@ class ProdForceSeAMaskOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == in_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == in_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); // Create an output tensor TensorShape force_shape; diff --git a/source/op/tf/prod_force_se_a_mask_grad.cc b/source/op/tf/prod_force_se_a_mask_grad.cc index c7ff091857..effb40759e 100644 --- a/source/op/tf/prod_force_se_a_mask_grad.cc +++ b/source/op/tf/prod_force_se_a_mask_grad.cc @@ -43,15 +43,17 @@ class ProdForceSeAMaskGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (mask_shape.dims() == 2), - errors::InvalidArgument("Dim of mask should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of mask should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); int nframes = net_deriv_tensor.shape().dim_size(0); int nloc = total_atom_num; @@ -59,22 +61,27 @@ class ProdForceSeAMaskGradOp : public OpKernel { int nnei = total_atom_num; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == mask_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (nloc * 3 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == mask_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_force_se_r_grad.cc b/source/op/tf/prod_force_se_r_grad.cc index 44741d20fb..9c6bd756fe 100644 --- a/source/op/tf/prod_force_se_r_grad.cc +++ b/source/op/tf/prod_force_se_r_grad.cc @@ -40,18 +40,21 @@ class ProdForceSeRGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -61,20 +64,24 @@ class ProdForceSeRGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (nloc * 3 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_virial.cc b/source/op/tf/prod_virial.cc index 42e7be669d..bf8fa94324 100644 --- a/source/op/tf/prod_virial.cc +++ b/source/op/tf/prod_virial.cc @@ -42,21 +42,24 @@ class ProdVirialOp : public OpKernel { const Tensor& natoms_tensor = context->input(5); // set size of the sample - OP_REQUIRES(context, (net_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + OP_REQUIRES( + context, (net_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); OP_REQUIRES(context, (axis_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of axis should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of axis should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -67,28 +70,35 @@ class ProdVirialOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == in_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == rij_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == axis_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == in_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == rij_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == axis_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 12 == in_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, - (static_cast(nloc) * nnei * 3 == - rij_tensor.shape().dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); + OP_REQUIRES( + context, + (static_cast(nloc) * nnei * 3 == + rij_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); OP_REQUIRES( - context, (nloc * 4 == axis_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of axis type+id should be 2+2")); + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); + OP_REQUIRES(context, (nloc * 4 == axis_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "number of axis type+id should be 2+2")); // Create an output tensor TensorShape virial_shape; diff --git a/source/op/tf/prod_virial_grad.cc b/source/op/tf/prod_virial_grad.cc index a764e524a6..7313bf5e29 100644 --- a/source/op/tf/prod_virial_grad.cc +++ b/source/op/tf/prod_virial_grad.cc @@ -49,22 +49,25 @@ class ProdVirialGradOp : public OpKernel { TensorShape axis_shape = axis_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_shape.dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); OP_REQUIRES(context, (axis_shape.dims() == 2), - errors::InvalidArgument("Dim of axis should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of axis should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -74,33 +77,40 @@ class ProdVirialGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == rij_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == axis_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (9 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == rij_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == axis_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (9 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 12 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); OP_REQUIRES( context, (static_cast(nloc) * nnei * 3 == rij_shape.dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); OP_REQUIRES( - context, (nloc * 4 == axis_shape.dim_size(1)), - errors::InvalidArgument("number of axis type+id should be 2+2")); + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); + OP_REQUIRES(context, (nloc * 4 == axis_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "number of axis type+id should be 2+2")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_virial_grad_multi_device.cc b/source/op/tf/prod_virial_grad_multi_device.cc index 3d8d2a96b3..fb011cabbd 100644 --- a/source/op/tf/prod_virial_grad_multi_device.cc +++ b/source/op/tf/prod_virial_grad_multi_device.cc @@ -57,20 +57,23 @@ class ProdVirialSeAGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_shape.dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -80,25 +83,32 @@ class ProdVirialSeAGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == rij_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (9 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == rij_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (9 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, (int_64(nloc) * nnei * 3 == rij_shape.dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); + OP_REQUIRES( + context, (int_64(nloc) * nnei * 3 == rij_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape grad_net_shape; @@ -191,20 +201,23 @@ class ProdVirialSeRGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_shape.dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -214,23 +227,29 @@ class ProdVirialSeRGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == rij_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (9 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == rij_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (9 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (int_64(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, (int_64(nloc) * nnei * 3 == rij_shape.dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); + OP_REQUIRES( + context, (int_64(nloc) * nnei * 3 == rij_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_virial_multi_device.cc b/source/op/tf/prod_virial_multi_device.cc index a544b010c5..5fff5cac8e 100644 --- a/source/op/tf/prod_virial_multi_device.cc +++ b/source/op/tf/prod_virial_multi_device.cc @@ -54,18 +54,21 @@ class ProdVirialSeAOp : public OpKernel { const Tensor& nlist_tensor = context->input(context_input_index++); const Tensor& natoms_tensor = context->input(context_input_index++); // set size of the sample - OP_REQUIRES(context, (net_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + OP_REQUIRES( + context, (net_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); const int* natoms = natoms_tensor.flat().data(); int nloc = natoms[0]; @@ -74,19 +77,23 @@ class ProdVirialSeAOp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == in_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == rij_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == in_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == rij_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); OP_REQUIRES( context, (int_64(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, - (int_64(nloc) * nnei * 3 == rij_tensor.shape().dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); + OP_REQUIRES( + context, (int_64(nloc) * nnei * 3 == rij_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); // Create an output tensor TensorShape virial_shape; virial_shape.AddDim(nframes); @@ -153,18 +160,21 @@ class ProdVirialSeROp : public OpKernel { const Tensor& nlist_tensor = context->input(context_input_index++); const Tensor& natoms_tensor = context->input(context_input_index++); // set size of the sample - OP_REQUIRES(context, (net_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + OP_REQUIRES( + context, (net_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); const int* natoms = natoms_tensor.flat().data(); int nloc = natoms[0]; @@ -173,19 +183,23 @@ class ProdVirialSeROp : public OpKernel { int nframes = net_deriv_tensor.shape().dim_size(0); int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == in_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == rij_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == in_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == rij_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); OP_REQUIRES( context, (int_64(nloc) * ndescrpt * 3 == in_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); - OP_REQUIRES(context, - (int_64(nloc) * nnei * 3 == rij_tensor.shape().dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); + OP_REQUIRES( + context, (int_64(nloc) * nnei * 3 == rij_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); // Create an output tensor TensorShape virial_shape; virial_shape.AddDim(nframes); diff --git a/source/op/tf/prod_virial_se_a_grad.cc b/source/op/tf/prod_virial_se_a_grad.cc index e3a9374b8f..c41a64f894 100644 --- a/source/op/tf/prod_virial_se_a_grad.cc +++ b/source/op/tf/prod_virial_se_a_grad.cc @@ -49,20 +49,23 @@ class ProdVirialSeAGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_shape.dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -72,28 +75,34 @@ class ProdVirialSeAGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == rij_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (9 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == rij_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (9 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); OP_REQUIRES( context, (static_cast(nloc) * nnei * 3 == rij_shape.dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/prod_virial_se_r_grad.cc b/source/op/tf/prod_virial_se_r_grad.cc index 8e9b2c25b0..55675c7104 100644 --- a/source/op/tf/prod_virial_se_r_grad.cc +++ b/source/op/tf/prod_virial_se_r_grad.cc @@ -43,20 +43,23 @@ class ProdVirialSeRGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (net_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (in_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (net_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (in_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_shape.dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -66,26 +69,31 @@ class ProdVirialSeRGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == in_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == rij_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES( - context, (9 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == in_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == rij_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES(context, (9 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES(context, (static_cast(nloc) * ndescrpt * 3 == in_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); OP_REQUIRES( context, (static_cast(nloc) * nnei * 3 == rij_shape.dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/soft_min.cc b/source/op/tf/soft_min.cc index 07c7404bbf..b239c166d6 100644 --- a/source/op/tf/soft_min.cc +++ b/source/op/tf/soft_min.cc @@ -52,16 +52,17 @@ class SoftMinSwitchOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (type_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of type should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of type should be 2")); OP_REQUIRES(context, (rij_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -73,23 +74,28 @@ class SoftMinSwitchOp : public OpKernel { assert(sel_r.size() == ntypes); // check the sizes - OP_REQUIRES(context, (nframes == type_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == rij_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nall == type_tensor.shape().dim_size(1)), - errors::InvalidArgument("shape of type should be nall")); OP_REQUIRES( - context, - (3 * static_cast(nnei) * nloc == - rij_tensor.shape().dim_size(1)), - errors::InvalidArgument("shape of rij should be 3 * nloc * nnei")); + context, (nframes == type_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == rij_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nall == type_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("shape of type should be nall")); + OP_REQUIRES(context, + (3 * static_cast(nnei) * nloc == + rij_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "shape of rij should be 3 * nloc * nnei")); OP_REQUIRES( context, (static_cast(nnei) * nloc == nlist_tensor.shape().dim_size(1)), - errors::InvalidArgument("shape of nlist should be nloc * nnei")); + deepmd::tf_compat::InvalidArgument( + "shape of nlist should be nloc * nnei")); // Create an output tensor TensorShape sw_value_shape; diff --git a/source/op/tf/soft_min_force.cc b/source/op/tf/soft_min_force.cc index 14cb42b993..5564141fb1 100644 --- a/source/op/tf/soft_min_force.cc +++ b/source/op/tf/soft_min_force.cc @@ -38,16 +38,18 @@ class SoftMinForceOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (du_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of du should be 2")); - OP_REQUIRES(context, (sw_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of switch deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of du should be 2")); + OP_REQUIRES( + context, (sw_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of switch deriv should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -57,19 +59,24 @@ class SoftMinForceOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == sw_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - - OP_REQUIRES(context, (nloc == du_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of du should match")); + OP_REQUIRES( + context, (nframes == sw_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + + OP_REQUIRES( + context, (nloc == du_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of du should match")); OP_REQUIRES(context, (static_cast(nloc) * nnei * 3 == sw_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of switch deriv should match")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument( + "number of switch deriv should match")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape force_shape; diff --git a/source/op/tf/soft_min_force_grad.cc b/source/op/tf/soft_min_force_grad.cc index e173586d22..08b3ab2da2 100644 --- a/source/op/tf/soft_min_force_grad.cc +++ b/source/op/tf/soft_min_force_grad.cc @@ -45,18 +45,20 @@ class SoftMinForceGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); OP_REQUIRES(context, (du_shape.dims() == 2), - errors::InvalidArgument("Dim of du should be 2")); - OP_REQUIRES(context, (sw_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of sw deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of du should be 2")); + OP_REQUIRES( + context, (sw_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of sw deriv should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -65,24 +67,29 @@ class SoftMinForceGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == sw_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - - OP_REQUIRES(context, (nloc == du_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of du should match")); OP_REQUIRES( - context, (nloc * 3 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == sw_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + + OP_REQUIRES( + context, (nloc == du_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of du should match")); + OP_REQUIRES(context, (nloc * 3 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES( context, (static_cast(nloc) * nnei * 3 == sw_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of sw deriv should match")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument("number of sw deriv should match")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/soft_min_virial.cc b/source/op/tf/soft_min_virial.cc index 6c0e1f72f3..40591d4cf6 100644 --- a/source/op/tf/soft_min_virial.cc +++ b/source/op/tf/soft_min_virial.cc @@ -41,19 +41,22 @@ class SoftMinVirialOp : public OpKernel { const Tensor& natoms_tensor = context->input(context_input_index++); // set size of the sample - OP_REQUIRES(context, (du_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (sw_deriv_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + OP_REQUIRES( + context, (du_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (sw_deriv_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -63,25 +66,32 @@ class SoftMinVirialOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == sw_deriv_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == rij_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - OP_REQUIRES(context, (nframes == nlist_tensor.shape().dim_size(0)), - errors::InvalidArgument("number of samples should match")); - - OP_REQUIRES(context, (nloc == du_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of du should match")); - OP_REQUIRES(context, - (static_cast(nloc) * nnei * 3 == - sw_deriv_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of sw_deriv should match")); - OP_REQUIRES(context, - (static_cast(nloc) * nnei * 3 == - rij_tensor.shape().dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + OP_REQUIRES( + context, (nframes == sw_deriv_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == rij_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + OP_REQUIRES( + context, (nframes == nlist_tensor.shape().dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of samples should match")); + + OP_REQUIRES( + context, (nloc == du_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of du should match")); + OP_REQUIRES( + context, + (static_cast(nloc) * nnei * 3 == + sw_deriv_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of sw_deriv should match")); + OP_REQUIRES( + context, + (static_cast(nloc) * nnei * 3 == + rij_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape virial_shape; diff --git a/source/op/tf/soft_min_virial_grad.cc b/source/op/tf/soft_min_virial_grad.cc index ac129b29af..4cc0c777be 100644 --- a/source/op/tf/soft_min_virial_grad.cc +++ b/source/op/tf/soft_min_virial_grad.cc @@ -48,20 +48,23 @@ class SoftMinVirialGradOp : public OpKernel { TensorShape nlist_shape = nlist_tensor.shape(); OP_REQUIRES(context, (grad_shape.dims() == 2), - errors::InvalidArgument("Dim of grad should be 2")); - OP_REQUIRES(context, (du_shape.dims() == 2), - errors::InvalidArgument("Dim of net deriv should be 2")); - OP_REQUIRES(context, (sw_deriv_shape.dims() == 2), - errors::InvalidArgument("Dim of input deriv should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of grad should be 2")); + OP_REQUIRES( + context, (du_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of net deriv should be 2")); + OP_REQUIRES( + context, (sw_deriv_shape.dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of input deriv should be 2")); OP_REQUIRES(context, (rij_shape.dims() == 2), - errors::InvalidArgument("Dim of rij should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of rij should be 2")); OP_REQUIRES(context, (nlist_shape.dims() == 2), - errors::InvalidArgument("Dim of nlist should be 2")); - OP_REQUIRES(context, (natoms_tensor.shape().dims() == 1), - errors::InvalidArgument("Dim of natoms should be 1")); + deepmd::tf_compat::InvalidArgument("Dim of nlist should be 2")); + OP_REQUIRES( + context, (natoms_tensor.shape().dims() == 1), + deepmd::tf_compat::InvalidArgument("Dim of natoms should be 1")); OP_REQUIRES(context, (natoms_tensor.shape().dim_size(0) >= 3), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "number of atoms should be larger than (or equal to) 3")); auto natoms = natoms_tensor.flat(); @@ -70,30 +73,37 @@ class SoftMinVirialGradOp : public OpKernel { int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0; // check the sizes - OP_REQUIRES(context, (nframes == grad_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == sw_deriv_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == rij_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); - OP_REQUIRES(context, (nframes == nlist_shape.dim_size(0)), - errors::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == grad_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == sw_deriv_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == rij_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES( + context, (nframes == nlist_shape.dim_size(0)), + deepmd::tf_compat::InvalidArgument("number of frames should match")); + OP_REQUIRES(context, (9 == grad_shape.dim_size(1)), + deepmd::tf_compat::InvalidArgument( + "input grad shape should be 3 x natoms")); OP_REQUIRES( - context, (9 == grad_shape.dim_size(1)), - errors::InvalidArgument("input grad shape should be 3 x natoms")); - OP_REQUIRES(context, (nloc == du_tensor.shape().dim_size(1)), - errors::InvalidArgument("number of du should match")); + context, (nloc == du_tensor.shape().dim_size(1)), + deepmd::tf_compat::InvalidArgument("number of du should match")); OP_REQUIRES( context, (static_cast(nloc) * nnei * 3 == sw_deriv_shape.dim_size(1)), - errors::InvalidArgument("number of descriptors should match")); + deepmd::tf_compat::InvalidArgument( + "number of descriptors should match")); OP_REQUIRES( context, (static_cast(nloc) * nnei * 3 == rij_shape.dim_size(1)), - errors::InvalidArgument("dim of rij should be nnei * 3")); - OP_REQUIRES(context, (nnei == n_a_sel + n_r_sel), - errors::InvalidArgument("number of neighbors should match")); + deepmd::tf_compat::InvalidArgument("dim of rij should be nnei * 3")); + OP_REQUIRES( + context, (nnei == n_a_sel + n_r_sel), + deepmd::tf_compat::InvalidArgument("number of neighbors should match")); // Create an output tensor TensorShape grad_net_shape; diff --git a/source/op/tf/tabulate_multi_device.cc b/source/op/tf/tabulate_multi_device.cc index 50267df556..174faf5213 100644 --- a/source/op/tf/tabulate_multi_device.cc +++ b/source/op/tf/tabulate_multi_device.cc @@ -184,11 +184,11 @@ class TabulateFusionSeAOp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of table should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); OP_REQUIRES(context, (em_x_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (em_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of input should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); TensorShape descriptor_shape; descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); descriptor_shape.AddDim(4); // be careful here; @@ -247,7 +247,7 @@ class TabulateFusionSeAGradOp : public OpKernel { const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of table should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); int context_output_index = 0; Tensor* dy_dem_x_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -311,9 +311,9 @@ class TabulateFusionSeAGradGradOp : public OpKernel { const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of input should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); int context_output_index = 0; Tensor* dz_dy_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -342,7 +342,7 @@ class TabulateFusionSeAGradGradOp : public OpKernel { dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM OP_REQUIRES(context, (last_layer_size <= 1024), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "In the process of model compression, the size of the " "last layer of embedding net must be less than 1024!")); } else if (device == "CPU") { @@ -381,13 +381,13 @@ class TabulateFusionSeAttenOp : public OpKernel { const Tensor& two_embed_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of table should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); OP_REQUIRES(context, (em_x_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (em_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of input should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); OP_REQUIRES(context, (two_embed_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); TensorShape descriptor_shape; descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); descriptor_shape.AddDim(4); // be careful here; @@ -452,7 +452,7 @@ class TabulateFusionSeAttenGradOp : public OpKernel { const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of table should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); int context_output_index = 0; Tensor* dy_dem_x_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -528,9 +528,9 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel { const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of input should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); int context_output_index = 0; Tensor* dz_dy_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -559,7 +559,7 @@ class TabulateFusionSeAttenGradGradOp : public OpKernel { dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM OP_REQUIRES(context, (last_layer_size <= 1024), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "In the process of model compression, the size of the " "last layer of embedding net must be less than 1024!")); } else if (device == "CPU") { @@ -596,11 +596,13 @@ class TabulateFusionSeTOp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of table should be 2")); - OP_REQUIRES(context, (em_x_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of em_x_tensor should be 2")); - OP_REQUIRES(context, (em_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of em_tensor should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); + OP_REQUIRES( + context, (em_x_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of em_x_tensor should be 2")); + OP_REQUIRES( + context, (em_tensor.shape().dims() == 3), + deepmd::tf_compat::InvalidArgument("Dim of em_tensor should be 3")); TensorShape descriptor_shape; descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); descriptor_shape.AddDim(last_layer_size); @@ -656,8 +658,9 @@ class TabulateFusionSeTGradOp : public OpKernel { const Tensor& dy_tensor = context->input(context_input_index++); const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample - OP_REQUIRES(context, (dy_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of dy_tensor should be 2")); + OP_REQUIRES( + context, (dy_tensor.shape().dims() == 2), + deepmd::tf_compat::InvalidArgument("Dim of dy_tensor should be 2")); int context_output_index = 0; Tensor* dy_dem_x_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -717,9 +720,9 @@ class TabulateFusionSeTGradGradOp : public OpKernel { const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_x_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of input should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 3")); int context_output_index = 0; Tensor* dz_dy_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -747,7 +750,7 @@ class TabulateFusionSeTGradGradOp : public OpKernel { nnei_i, nnei_j, last_layer_size); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM OP_REQUIRES(context, (last_layer_size <= 1024), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "In the process of model compression, the size of the " "last layer of embedding net must be less than 1024!")); } else if (device == "CPU") { @@ -781,9 +784,9 @@ class TabulateFusionSeROp : public OpKernel { const Tensor& em_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (table_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of table should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of table should be 2")); OP_REQUIRES(context, (em_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); TensorShape descriptor_shape; descriptor_shape.AddDim(em_tensor.shape().dim_size(0)); descriptor_shape.AddDim(em_tensor.shape().dim_size(1)); // be careful here; @@ -838,7 +841,7 @@ class TabulateFusionSeRGradOp : public OpKernel { const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (dy_tensor.shape().dims() == 3), - errors::InvalidArgument("Dim of table should be 3")); + deepmd::tf_compat::InvalidArgument("Dim of table should be 3")); int context_output_index = 0; Tensor* dy_dem_tensor = NULL; OP_REQUIRES_OK(context, @@ -887,7 +890,7 @@ class TabulateFusionSeRGradGradOp : public OpKernel { const Tensor& descriptor_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (dz_dy_dem_tensor.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); int context_output_index = 0; Tensor* dz_dy_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, @@ -911,7 +914,7 @@ class TabulateFusionSeRGradGradOp : public OpKernel { dz_dy, table, table_info, em, dz_dy_dem, nloc, nnei, last_layer_size); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM OP_REQUIRES(context, (last_layer_size <= 1024), - errors::InvalidArgument( + deepmd::tf_compat::InvalidArgument( "In the process of model compression, the size of the " "last layer of embedding net must be less than 1024!")); } else if (device == "CPU") { diff --git a/source/op/tf/unaggregated_grad.cc b/source/op/tf/unaggregated_grad.cc index 329e25b2d2..ce74c5f438 100644 --- a/source/op/tf/unaggregated_grad.cc +++ b/source/op/tf/unaggregated_grad.cc @@ -314,11 +314,11 @@ class UnaggregatedDyDxSOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (y.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (w.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (xbar.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); // check functype int context_output_index = 0; @@ -360,13 +360,13 @@ class UnaggregatedDy2DxSOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (y.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dy.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (w.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (xbar.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); int context_output_index = 0; Tensor* dy2_dx = NULL; @@ -407,13 +407,13 @@ class UnaggregatedDyDxOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (z.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (w.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dy_dx.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (ybar.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); int context_output_index = 0; Tensor* dz_dx = NULL; @@ -457,15 +457,15 @@ class UnaggregatedDy2DxOp : public OpKernel { // set size of the sample OP_REQUIRES(context, (z.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (w.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dy_dx.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (dy2_dx.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); OP_REQUIRES(context, (ybar.shape().dims() == 2), - errors::InvalidArgument("Dim of input should be 2")); + deepmd::tf_compat::InvalidArgument("Dim of input should be 2")); int context_output_index = 0; Tensor* dz2_dx = NULL;