# OP API

## op_module

Python wrappers around TensorFlow ops.

This file is MACHINE GENERATED! Do not edit.

deepmd.env.op_module.Descrpt(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, sel_a, sel_r, axis_rule, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• axis_rule – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist, axis, rot_mat).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32. axis: A Tensor of type int32. rot_mat: A Tensor. Has the same type as coord.

deepmd.env.op_module.DescrptNorot(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.DescrptSeA(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.DescrptSeAEf(coord, type, natoms, box, mesh, ef, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• ef – A Tensor. Must have the same type as coord.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.DescrptSeAEfPara(coord, type, natoms, box, mesh, ef, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• ef – A Tensor. Must have the same type as coord.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.DescrptSeAEfVert(coord, type, natoms, box, mesh, ef, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• ef – A Tensor. Must have the same type as coord.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.DescrptSeR(coord, type, natoms, box, mesh, davg, dstd, rcut, rcut_smth, sel, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut – A float.

• rcut_smth – A float.

• sel – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.EwaldRecp(coord, charge, natoms, box, ewald_beta, ewald_h, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• charge – A Tensor. Must have the same type as coord.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• ewald_beta – A float.

• ewald_h – A float.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (energy, force, virial).

energy: A Tensor. Has the same type as coord. force: A Tensor. Has the same type as coord. virial: A Tensor. Has the same type as coord.

deepmd.env.op_module.Gelu(x, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.GeluCustom(x, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

Parameters
• dy – A Tensor. Must be one of the following types: float32, float64.

• x – A Tensor. Must have the same type as dy.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as dy.

Parameters
• dy – A Tensor. Must be one of the following types: float32, float64.

• x – A Tensor. Must have the same type as dy.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as dy.

Parameters
• dy – A Tensor. Must be one of the following types: float32, float64.

• dy – A Tensor. Must have the same type as dy.

• x – A Tensor. Must have the same type as dy.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as dy.

Parameters
• dy – A Tensor. Must be one of the following types: float32, float64.

• dy – A Tensor. Must have the same type as dy.

• x – A Tensor. Must have the same type as dy.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as dy.

deepmd.env.op_module.MapAparam(aparam, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• aparam – A Tensor. Must be one of the following types: float32, float64.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as aparam.

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• v – A Tensor. Must have the same type as x.

• dv – A Tensor. Must have the same type as x.

• grad_v – A Tensor. Must have the same type as x.

• grad_dv – A Tensor. Must have the same type as x.

• prec – A float.

• nbit – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.MatmulNvnmd(x, w, isround, nbit1, nbit2, nbit3, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• w – A Tensor. Must have the same type as x.

• isround – An int.

• nbit1 – An int.

• nbit2 – An int.

• nbit3 – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.NeighborStat(coord, type, natoms, box, mesh, rcut, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• rcut – A float.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (max_nbor_size, min_nbor_dist).

max_nbor_size: A Tensor of type int32. min_nbor_dist: A Tensor. Has the same type as coord.

deepmd.env.op_module.PairTab(table_info, table_data, type, rij, nlist, natoms, scale, sel_a, sel_r, name=None)

Parameters
• table_info – A Tensor of type float64.

• table_data – A Tensor of type float64.

• type – A Tensor of type int32.

• rij – A Tensor. Must be one of the following types: float32, float64.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• scale – A Tensor. Must have the same type as rij.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (atom_energy, force, atom_virial).

atom_energy: A Tensor. Has the same type as rij. force: A Tensor. Has the same type as rij. atom_virial: A Tensor. Has the same type as rij.

deepmd.env.op_module.ParallelProdForceSeA(net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, parallel=False, start_frac=0, end_frac=1, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• parallel – An optional bool. Defaults to False.

• start_frac – An optional float. Defaults to 0.

• end_frac – An optional float. Defaults to 1.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.ProdEnvMatA(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Compute the environment matrix for descriptor se_e2_a.

Each row of the environment matrix $$\mathcal{R}^i$$ can be constructed as follows

$(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]$

In the above equation, $$\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})$$ is the relative coordinate and $$r_{ji}=\lVert \mathbf{R}_{ji} \lVert$$ is its norm. The switching function $$s(r)$$ is defined as:

$\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}$

Note that the environment matrix is normalized by davg and dstd.

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64. The coordinates of atoms.

• type – A Tensor of type int32. The types of atoms.

• natoms – A Tensor of type int32. The number of atoms. This tensor has the length of Ntypes + 2. natoms[0]: number of local atoms. natoms[1]: total number of atoms held by this processor. natoms[i]: 2 <= i < Ntypes+2, number of type i atoms.

• box – A Tensor. Must have the same type as coord. The box of frames.

• mesh – A Tensor of type int32. Gor historical reasons, only the length of the Tensor matters. If size of mesh == 6, pbc is assumed. If size of mesh == 0, no-pbc is assumed.

• davg – A Tensor. Must have the same type as coord. Average value of the environment matrix for normalization.

• dstd – A Tensor. Must have the same type as coord. Standard deviation of the environment matrix for normalization.

• rcut_a – A float. This argument is not used.

• rcut_r – A float. The cutoff radius for the environment matrix.

• rcut_r_smth – A float. From where the environment matrix should be smoothed.

• sel_a – A list of ints. sel_a[i] specifies the maxmum number of type i atoms in the cut-off radius.

• sel_r – A list of ints. This argument is not used.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. The environment matrix. descrpt_deriv: A Tensor. Has the same type as coord. The derivative of the environment matrix. rij: A Tensor. Has the same type as coord. The distance between the atoms. nlist: A Tensor of type int32. The neighbor list of each atom.

deepmd.env.op_module.ProdEnvMatAMix(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Compute the environment matrix mixing the atom types.

The sorting of neighbor atoms depends not on atom types, but on the distance and index. The atoms in nlist matrix will gather forward and thus save space for gaps of types in ProdEnvMatA, resulting in optimized and relative small sel_a.

The additional outputs are listed as following:

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist, ntype, nmask).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32. ntype: A Tensor of type int32. The corresponding atom types in nlist. nmask: A Tensor of type bool. The atom mask in nlist.

deepmd.env.op_module.ProdEnvMatANvnmdQuantize(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.ProdEnvMatR(coord, type, natoms, box, mesh, davg, dstd, rcut, rcut_smth, sel, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut – A float.

• rcut_smth – A float.

• sel – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.ProdForce(net_deriv, in_deriv, nlist, axis, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• axis – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.ProdForceNorot(net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.ProdForceSeA(net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.ProdForceSeR(net_deriv, in_deriv, nlist, natoms, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.ProdVirial(net_deriv, in_deriv, rij, nlist, axis, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• rij – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• axis – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as net_deriv. atom_virial: A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.ProdVirialNorot(net_deriv, in_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• rij – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as net_deriv. atom_virial: A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.ProdVirialSeA(net_deriv, in_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• rij – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as net_deriv. atom_virial: A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.ProdVirialSeR(net_deriv, in_deriv, rij, nlist, natoms, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• rij – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as net_deriv. atom_virial: A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.QuantizeNvnmd(x, isround, nbit1, nbit2, nbit3, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• isround – An int.

• nbit1 – An int.

• nbit2 – An int.

• nbit3 – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.SoftMinForce(du, sw_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• du – A Tensor. Must be one of the following types: float32, float64.

• sw_deriv – A Tensor. Must have the same type as du.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as du.

deepmd.env.op_module.SoftMinSwitch(type, rij, nlist, natoms, sel_a, sel_r, alpha, rmin, rmax, name=None)

Parameters
• type – A Tensor of type int32.

• rij – A Tensor. Must be one of the following types: float32, float64.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• sel_a – A list of ints.

• sel_r – A list of ints.

• alpha – A float.

• rmin – A float.

• rmax – A float.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (sw_value, sw_deriv).

sw_value: A Tensor. Has the same type as rij. sw_deriv: A Tensor. Has the same type as rij.

deepmd.env.op_module.SoftMinVirial(du, sw_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• du – A Tensor. Must be one of the following types: float32, float64.

• sw_deriv – A Tensor. Must have the same type as du.

• rij – A Tensor. Must have the same type as du.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as du. atom_virial: A Tensor. Has the same type as du.

deepmd.env.op_module.TabulateFusion(table, table_info, em_x, em, last_layer_size, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• last_layer_size – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.TabulateFusionGrad(table, table_info, em_x, em, dy, descriptor, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dy – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (dy_dem_x, dy_dem).

dy_dem_x: A Tensor. Has the same type as table. dy_dem: A Tensor. Has the same type as table.

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dz_dy_dem_x – A Tensor. Must have the same type as table.

• dz_dy_dem – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.TabulateFusionSeA(table, table_info, em_x, em, last_layer_size, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• last_layer_size – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.TabulateFusionSeAGrad(table, table_info, em_x, em, dy, descriptor, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dy – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (dy_dem_x, dy_dem).

dy_dem_x: A Tensor. Has the same type as table. dy_dem: A Tensor. Has the same type as table.

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dz_dy_dem_x – A Tensor. Must have the same type as table.

• dz_dy_dem – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.TabulateFusionSeR(table, table_info, em, last_layer_size, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• last_layer_size – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.TabulateFusionSeRGrad(table, table_info, em, dy, descriptor, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dy – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dz_dy_dem – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.TabulateFusionSeT(table, table_info, em_x, em, last_layer_size, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• last_layer_size – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.TabulateFusionSeTGrad(table, table_info, em_x, em, dy, descriptor, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dy – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (dy_dem_x, dy_dem).

dy_dem_x: A Tensor. Has the same type as table. dy_dem: A Tensor. Has the same type as table.

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dz_dy_dem_x – A Tensor. Must have the same type as table.

• dz_dy_dem – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.Tanh2Nvnmd(x, isround, nbit1, nbit2, nbit3, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• isround – An int.

• nbit1 – An int.

• nbit2 – An int.

• nbit3 – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.Tanh4Nvnmd(x, isround, nbit1, nbit2, nbit3, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• isround – An int.

• nbit1 – An int.

• nbit2 – An int.

• nbit3 – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.UnaggregatedDy2Dx(z, w, dy_dx, dy2_dx, ybar, functype, name=None)

Parameters
• z – A Tensor. Must be one of the following types: float32, float64.

• w – A Tensor. Must have the same type as z.

• dy_dx – A Tensor. Must have the same type as z.

• dy2_dx – A Tensor. Must have the same type as z.

• ybar – A Tensor. Must have the same type as z.

• functype – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as z.

deepmd.env.op_module.UnaggregatedDy2DxS(y, dy, w, xbar, functype, name=None)

Parameters
• y – A Tensor. Must be one of the following types: float32, float64.

• dy – A Tensor. Must have the same type as y.

• w – A Tensor. Must have the same type as y.

• xbar – A Tensor. Must have the same type as y.

• functype – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as y.

deepmd.env.op_module.UnaggregatedDyDx(z, w, dy_dx, ybar, functype, name=None)

Parameters
• z – A Tensor. Must be one of the following types: float32, float64.

• w – A Tensor. Must have the same type as z.

• dy_dx – A Tensor. Must have the same type as z.

• ybar – A Tensor. Must have the same type as z.

• functype – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as z.

deepmd.env.op_module.UnaggregatedDyDxS(y, w, xbar, functype, name=None)

Parameters
• y – A Tensor. Must be one of the following types: float32, float64.

• w – A Tensor. Must have the same type as y.

• xbar – A Tensor. Must have the same type as y.

• functype – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as y.

deepmd.env.op_module.descrpt(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, sel_a, sel_r, axis_rule, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• axis_rule – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist, axis, rot_mat).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32. axis: A Tensor of type int32. rot_mat: A Tensor. Has the same type as coord.

deepmd.env.op_module.descrpt_norot(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.descrpt_se_a(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.descrpt_se_a_ef(coord, type, natoms, box, mesh, ef, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• ef – A Tensor. Must have the same type as coord.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.descrpt_se_a_ef_para(coord, type, natoms, box, mesh, ef, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• ef – A Tensor. Must have the same type as coord.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.descrpt_se_a_ef_vert(coord, type, natoms, box, mesh, ef, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• ef – A Tensor. Must have the same type as coord.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.descrpt_se_r(coord, type, natoms, box, mesh, davg, dstd, rcut, rcut_smth, sel, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut – A float.

• rcut_smth – A float.

• sel – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.ewald_recp(coord, charge, natoms, box, ewald_beta, ewald_h, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• charge – A Tensor. Must have the same type as coord.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• ewald_beta – A float.

• ewald_h – A float.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (energy, force, virial).

energy: A Tensor. Has the same type as coord. force: A Tensor. Has the same type as coord. virial: A Tensor. Has the same type as coord.

deepmd.env.op_module.gelu(x, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.gelu_custom(x, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

Parameters
• dy – A Tensor. Must be one of the following types: float32, float64.

• x – A Tensor. Must have the same type as dy.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as dy.

Parameters
• dy – A Tensor. Must be one of the following types: float32, float64.

• x – A Tensor. Must have the same type as dy.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as dy.

Parameters
• dy – A Tensor. Must be one of the following types: float32, float64.

• dy – A Tensor. Must have the same type as dy.

• x – A Tensor. Must have the same type as dy.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as dy.

Parameters
• dy – A Tensor. Must be one of the following types: float32, float64.

• dy – A Tensor. Must have the same type as dy.

• x – A Tensor. Must have the same type as dy.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as dy.

deepmd.env.op_module.map_aparam(aparam, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• aparam – A Tensor. Must be one of the following types: float32, float64.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as aparam.

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• v – A Tensor. Must have the same type as x.

• dv – A Tensor. Must have the same type as x.

• grad_v – A Tensor. Must have the same type as x.

• grad_dv – A Tensor. Must have the same type as x.

• prec – A float.

• nbit – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.matmul_nvnmd(x, w, isround, nbit1, nbit2, nbit3, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• w – A Tensor. Must have the same type as x.

• isround – An int.

• nbit1 – An int.

• nbit2 – An int.

• nbit3 – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.neighbor_stat(coord, type, natoms, box, mesh, rcut, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• rcut – A float.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (max_nbor_size, min_nbor_dist).

max_nbor_size: A Tensor of type int32. min_nbor_dist: A Tensor. Has the same type as coord.

deepmd.env.op_module.pair_tab(table_info, table_data, type, rij, nlist, natoms, scale, sel_a, sel_r, name=None)

Parameters
• table_info – A Tensor of type float64.

• table_data – A Tensor of type float64.

• type – A Tensor of type int32.

• rij – A Tensor. Must be one of the following types: float32, float64.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• scale – A Tensor. Must have the same type as rij.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (atom_energy, force, atom_virial).

atom_energy: A Tensor. Has the same type as rij. force: A Tensor. Has the same type as rij. atom_virial: A Tensor. Has the same type as rij.

deepmd.env.op_module.parallel_prod_force_se_a(net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, parallel=False, start_frac=0, end_frac=1, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• parallel – An optional bool. Defaults to False.

• start_frac – An optional float. Defaults to 0.

• end_frac – An optional float. Defaults to 1.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.prod_env_mat_a(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Compute the environment matrix for descriptor se_e2_a.

Each row of the environment matrix $$\mathcal{R}^i$$ can be constructed as follows

$(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]$

In the above equation, $$\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})$$ is the relative coordinate and $$r_{ji}=\lVert \mathbf{R}_{ji} \lVert$$ is its norm. The switching function $$s(r)$$ is defined as:

$\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}$

Note that the environment matrix is normalized by davg and dstd.

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64. The coordinates of atoms.

• type – A Tensor of type int32. The types of atoms.

• natoms – A Tensor of type int32. The number of atoms. This tensor has the length of Ntypes + 2. natoms[0]: number of local atoms. natoms[1]: total number of atoms held by this processor. natoms[i]: 2 <= i < Ntypes+2, number of type i atoms.

• box – A Tensor. Must have the same type as coord. The box of frames.

• mesh – A Tensor of type int32. Gor historical reasons, only the length of the Tensor matters. If size of mesh == 6, pbc is assumed. If size of mesh == 0, no-pbc is assumed.

• davg – A Tensor. Must have the same type as coord. Average value of the environment matrix for normalization.

• dstd – A Tensor. Must have the same type as coord. Standard deviation of the environment matrix for normalization.

• rcut_a – A float. This argument is not used.

• rcut_r – A float. The cutoff radius for the environment matrix.

• rcut_r_smth – A float. From where the environment matrix should be smoothed.

• sel_a – A list of ints. sel_a[i] specifies the maxmum number of type i atoms in the cut-off radius.

• sel_r – A list of ints. This argument is not used.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. The environment matrix. descrpt_deriv: A Tensor. Has the same type as coord. The derivative of the environment matrix. rij: A Tensor. Has the same type as coord. The distance between the atoms. nlist: A Tensor of type int32. The neighbor list of each atom.

deepmd.env.op_module.prod_env_mat_a_mix(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Compute the environment matrix mixing the atom types.

The sorting of neighbor atoms depends not on atom types, but on the distance and index. The atoms in nlist matrix will gather forward and thus save space for gaps of types in ProdEnvMatA, resulting in optimized and relative small sel_a.

The additional outputs are listed as following:

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist, ntype, nmask).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32. ntype: A Tensor of type int32. The corresponding atom types in nlist. nmask: A Tensor of type bool. The atom mask in nlist.

deepmd.env.op_module.prod_env_mat_a_nvnmd_quantize(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut_a – A float.

• rcut_r – A float.

• rcut_r_smth – A float.

• sel_a – A list of ints.

• sel_r – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.prod_env_mat_r(coord, type, natoms, box, mesh, davg, dstd, rcut, rcut_smth, sel, name=None)

Parameters
• coord – A Tensor. Must be one of the following types: float32, float64.

• type – A Tensor of type int32.

• natoms – A Tensor of type int32.

• box – A Tensor. Must have the same type as coord.

• mesh – A Tensor of type int32.

• davg – A Tensor. Must have the same type as coord.

• dstd – A Tensor. Must have the same type as coord.

• rcut – A float.

• rcut_smth – A float.

• sel – A list of ints.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (descrpt, descrpt_deriv, rij, nlist).

descrpt: A Tensor. Has the same type as coord. descrpt_deriv: A Tensor. Has the same type as coord. rij: A Tensor. Has the same type as coord. nlist: A Tensor of type int32.

deepmd.env.op_module.prod_force(net_deriv, in_deriv, nlist, axis, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• axis – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.prod_force_norot(net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.prod_force_se_a(net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.prod_force_se_r(net_deriv, in_deriv, nlist, natoms, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.prod_virial(net_deriv, in_deriv, rij, nlist, axis, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• rij – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• axis – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as net_deriv. atom_virial: A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.prod_virial_norot(net_deriv, in_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• rij – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as net_deriv. atom_virial: A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.prod_virial_se_a(net_deriv, in_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• rij – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as net_deriv. atom_virial: A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.prod_virial_se_r(net_deriv, in_deriv, rij, nlist, natoms, name=None)

Parameters
• net_deriv – A Tensor. Must be one of the following types: float32, float64.

• in_deriv – A Tensor. Must have the same type as net_deriv.

• rij – A Tensor. Must have the same type as net_deriv.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as net_deriv. atom_virial: A Tensor. Has the same type as net_deriv.

deepmd.env.op_module.quantize_nvnmd(x, isround, nbit1, nbit2, nbit3, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• isround – An int.

• nbit1 – An int.

• nbit2 – An int.

• nbit3 – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.soft_min_force(du, sw_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• du – A Tensor. Must be one of the following types: float32, float64.

• sw_deriv – A Tensor. Must have the same type as du.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as du.

deepmd.env.op_module.soft_min_switch(type, rij, nlist, natoms, sel_a, sel_r, alpha, rmin, rmax, name=None)

Parameters
• type – A Tensor of type int32.

• rij – A Tensor. Must be one of the following types: float32, float64.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• sel_a – A list of ints.

• sel_r – A list of ints.

• alpha – A float.

• rmin – A float.

• rmax – A float.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (sw_value, sw_deriv).

sw_value: A Tensor. Has the same type as rij. sw_deriv: A Tensor. Has the same type as rij.

deepmd.env.op_module.soft_min_virial(du, sw_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

Parameters
• du – A Tensor. Must be one of the following types: float32, float64.

• sw_deriv – A Tensor. Must have the same type as du.

• rij – A Tensor. Must have the same type as du.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (virial, atom_virial).

virial: A Tensor. Has the same type as du. atom_virial: A Tensor. Has the same type as du.

deepmd.env.op_module.tabulate_fusion(table, table_info, em_x, em, last_layer_size, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• last_layer_size – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.tabulate_fusion_grad(table, table_info, em_x, em, dy, descriptor, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dy – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (dy_dem_x, dy_dem).

dy_dem_x: A Tensor. Has the same type as table. dy_dem: A Tensor. Has the same type as table.

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dz_dy_dem_x – A Tensor. Must have the same type as table.

• dz_dy_dem – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.tabulate_fusion_se_a(table, table_info, em_x, em, last_layer_size, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• last_layer_size – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.tabulate_fusion_se_a_grad(table, table_info, em_x, em, dy, descriptor, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dy – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (dy_dem_x, dy_dem).

dy_dem_x: A Tensor. Has the same type as table. dy_dem: A Tensor. Has the same type as table.

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dz_dy_dem_x – A Tensor. Must have the same type as table.

• dz_dy_dem – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.tabulate_fusion_se_r(table, table_info, em, last_layer_size, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• last_layer_size – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.tabulate_fusion_se_r_grad(table, table_info, em, dy, descriptor, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dy – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dz_dy_dem – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.tabulate_fusion_se_t(table, table_info, em_x, em, last_layer_size, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• last_layer_size – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.tabulate_fusion_se_t_grad(table, table_info, em_x, em, dy, descriptor, name=None)

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dy – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A tuple of Tensor objects (dy_dem_x, dy_dem).

dy_dem_x: A Tensor. Has the same type as table. dy_dem: A Tensor. Has the same type as table.

Parameters
• table – A Tensor. Must be one of the following types: float32, float64.

• table_info – A Tensor. Must have the same type as table.

• em_x – A Tensor. Must have the same type as table.

• em – A Tensor. Must have the same type as table.

• dz_dy_dem_x – A Tensor. Must have the same type as table.

• dz_dy_dem – A Tensor. Must have the same type as table.

• descriptor – A Tensor. Must have the same type as table.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as table.

deepmd.env.op_module.tanh2_nvnmd(x, isround, nbit1, nbit2, nbit3, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• isround – An int.

• nbit1 – An int.

• nbit2 – An int.

• nbit3 – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.tanh4_nvnmd(x, isround, nbit1, nbit2, nbit3, name=None)

Parameters
• x – A Tensor. Must be one of the following types: float32, float64.

• isround – An int.

• nbit1 – An int.

• nbit2 – An int.

• nbit3 – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as x.

deepmd.env.op_module.unaggregated_dy2_dx(z, w, dy_dx, dy2_dx, ybar, functype, name=None)

Parameters
• z – A Tensor. Must be one of the following types: float32, float64.

• w – A Tensor. Must have the same type as z.

• dy_dx – A Tensor. Must have the same type as z.

• dy2_dx – A Tensor. Must have the same type as z.

• ybar – A Tensor. Must have the same type as z.

• functype – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as z.

deepmd.env.op_module.unaggregated_dy2_dx_s(y, dy, w, xbar, functype, name=None)

Parameters
• y – A Tensor. Must be one of the following types: float32, float64.

• dy – A Tensor. Must have the same type as y.

• w – A Tensor. Must have the same type as y.

• xbar – A Tensor. Must have the same type as y.

• functype – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as y.

deepmd.env.op_module.unaggregated_dy_dx(z, w, dy_dx, ybar, functype, name=None)

Parameters
• z – A Tensor. Must be one of the following types: float32, float64.

• w – A Tensor. Must have the same type as z.

• dy_dx – A Tensor. Must have the same type as z.

• ybar – A Tensor. Must have the same type as z.

• functype – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as z.

deepmd.env.op_module.unaggregated_dy_dx_s(y, w, xbar, functype, name=None)

Parameters
• y – A Tensor. Must be one of the following types: float32, float64.

• w – A Tensor. Must have the same type as y.

• xbar – A Tensor. Must have the same type as y.

• functype – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as y.

Python wrappers around TensorFlow ops.

This file is MACHINE GENERATED! Do not edit.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• axis – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• rij – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• axis – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• rij – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• rij – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• du – A Tensor. Must have the same type as grad.

• sw_deriv – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• du – A Tensor. Must have the same type as grad.

• sw_deriv – A Tensor. Must have the same type as grad.

• rij – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• axis – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• rij – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• axis – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• rij – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• net_deriv – A Tensor. Must have the same type as grad.

• in_deriv – A Tensor. Must have the same type as grad.

• rij – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• du – A Tensor. Must have the same type as grad.

• sw_deriv – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.

Parameters
• grad – A Tensor. Must be one of the following types: float32, float64.

• du – A Tensor. Must have the same type as grad.

• sw_deriv – A Tensor. Must have the same type as grad.

• rij – A Tensor. Must have the same type as grad.

• nlist – A Tensor of type int32.

• natoms – A Tensor of type int32.

• n_a_sel – An int.

• n_r_sel – An int.

• name – A name for the operation (optional).

Returns

A Tensor. Has the same type as grad.