OP API

op_module

Python wrappers around TensorFlow ops.

This file is MACHINE GENERATED! Do not edit.

deepmd.tf.env.op_module.AddFltNvnmd(x, w, name=None)

TODO: add doc.

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

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

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

Returns

A Tensor. Has the same type as x.

class deepmd.tf.env.op_module.Annotated(*args, **kwargs)

Bases: object

Add context specific metadata to a type.

Example: Annotated[int, runtime_check.Unsigned] indicates to the hypothetical runtime_check module that this type is an unsigned int. Every other consumer of this type can ignore this metadata and treat this type as int.

The first argument to Annotated must be a valid type.

Details:

  • It’s an error to call Annotated with less than two arguments.

  • Nested Annotated are flattened:

    Annotated[Annotated[T, Ann1, Ann2], Ann3] == Annotated[T, Ann1, Ann2, Ann3]
    
  • Instantiating an annotated type is equivalent to instantiating the

underlying type:

Annotated[C, Ann1](5) == C(5)
  • Annotated can be used as a generic type alias:

    Optimized = Annotated[T, runtime.Optimize()]
    Optimized[int] == Annotated[int, runtime.Optimize()]
    
    OptimizedList = Annotated[List[T], runtime.Optimize()]
    OptimizedList[int] == Annotated[List[int], runtime.Optimize()]
    
deepmd.tf.env.op_module.ConvertForwardMap(sub_forward_map, sub_natoms, natoms, name=None)

TODO: add doc.

Parameters
  • sub_forward_map – A Tensor of type int32.

  • sub_natoms – A Tensor of type int32.

  • natoms – A Tensor of type int32.

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

Returns

A tuple of Tensor objects (forward_map, backward_map, new_natoms, mesh).

forward_map: A Tensor of type int32. backward_map: A Tensor of type int32. new_natoms: A Tensor of type int32. mesh: A Tensor of type int32.

deepmd.tf.env.op_module.CopyFltNvnmd(x, name=None)

TODO: add doc.

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

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

Returns

A tuple of Tensor objects (y1, y2).

y1: A Tensor. Has the same type as x. y2: A Tensor. Has the same type as x.

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

TODO: add doc.

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.tf.env.op_module.DescrptNorot(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.env.op_module.DescrptSeA(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.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)

TODO: add doc.

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.tf.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)

TODO: add doc.

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.tf.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)

TODO: add doc.

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.tf.env.op_module.DescrptSeAMask(coord, type, mask, box, natoms, mesh, name=None)

TODO: add doc.

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

  • type – A Tensor of type int32.

  • mask – A Tensor of type int32.

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

  • natoms – A Tensor of type int32.

  • mesh – A Tensor of type int32.

  • 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.tf.env.op_module.DescrptSeR(coord, type, natoms, box, mesh, davg, dstd, rcut, rcut_smth, sel, name=None)

TODO: add doc.

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.tf.env.op_module.DotmulFltNvnmd(x, w, name=None)

TODO: add doc.

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

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

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.DprcPairwiseIdx(idxs, natoms, name=None)

TODO: add doc.

Parameters
  • idxs – A Tensor of type int32.

  • natoms – A Tensor of type int32.

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

Returns

A tuple of Tensor objects (forward_qm_map, backward_qm_map, forward_qmmm_map, backward_qmmm_map, natoms_qm, natoms_qmmm, qmmm_frame_idx).

forward_qm_map: A Tensor of type int32. backward_qm_map: A Tensor of type int32. forward_qmmm_map: A Tensor of type int32. backward_qmmm_map: A Tensor of type int32. natoms_qm: A Tensor of type int32. natoms_qmmm: A Tensor of type int32. qmmm_frame_idx: A Tensor of type int32.

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

TODO: add doc.

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.tf.env.op_module.FltNvnmd(x, name=None)

TODO: add doc.

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.tf.env.op_module.Gelu(x, name=None)

TODO: add doc.

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.tf.env.op_module.GeluCustom(x, name=None)

TODO: add doc.

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.tf.env.op_module.GeluGrad(dy, x, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.GeluGradCustom(dy, x, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.GeluGradGrad(dy, dy_, x, name=None)

TODO: add doc.

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.tf.env.op_module.GeluGradGradCustom(dy, dy_, x, name=None)

TODO: add doc.

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.tf.env.op_module.MapAparam(aparam, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.MapFltNvnmd(x, table, table_grad, table_info, name=None)

TODO: add doc.

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

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

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

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

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.MatmulFitnetNvnmd(x, w, nbitx, nbitw, normw, name=None)

TODO: add doc.

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

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

  • nbitx – An int.

  • nbitw – An int.

  • normw – An int.

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.MatmulFlt2fixNvnmd(x, w, nbit, name=None)

TODO: add doc.

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

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

  • nbit – An int.

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.MatmulFltNvnmd(x, w, normx, normw, name=None)

TODO: add doc.

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

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

  • normx – An int.

  • normw – An int.

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.MulFltNvnmd(x, w, name=None)

TODO: add doc.

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

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

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

Returns

A Tensor. Has the same type as x.

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

TODO: add doc.

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.tf.env.op_module.PairTab(table_info, table_data, type, rij, nlist, natoms, scale, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.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)

TODO: add doc.

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.tf.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.tf.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.tf.env.op_module.ProdEnvMatAMixNvnmdQuantize(coord, type, natoms, box, mesh, davg, dstd, rcut_a, rcut_r, rcut_r_smth, sel_a, sel_r, name=None)

TODO: add doc.

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. nmask: A Tensor of type bool.

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

TODO: add doc.

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.tf.env.op_module.ProdEnvMatR(coord, type, natoms, box, mesh, davg, dstd, rcut, rcut_smth, sel, name=None)

TODO: add doc.

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.tf.env.op_module.ProdForce(net_deriv, in_deriv, nlist, axis, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.tf.env.op_module.ProdForceNorot(net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.tf.env.op_module.ProdForceSeA(net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.tf.env.op_module.ProdForceSeAMask(net_deriv, in_deriv, mask, nlist, total_atom_num, name=None)

TODO: add doc.

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.

  • mask – A Tensor of type int32.

  • nlist – A Tensor of type int32.

  • total_atom_num – An int.

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

Returns

A Tensor. Has the same type as net_deriv.

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

TODO: add doc.

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.tf.env.op_module.ProdVirial(net_deriv, in_deriv, rij, nlist, axis, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.tf.env.op_module.ProdVirialNorot(net_deriv, in_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.tf.env.op_module.ProdVirialSeA(net_deriv, in_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.tf.env.op_module.ProdVirialSeR(net_deriv, in_deriv, rij, nlist, natoms, name=None)

TODO: add doc.

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.tf.env.op_module.QuantizeNvnmd(x, isround, nbit1, nbit2, nbit3, name=None)

TODO: add doc.

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.tf.env.op_module.SoftMinForce(du, sw_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.tf.env.op_module.SoftMinSwitch(type, rij, nlist, natoms, sel_a, sel_r, alpha, rmin, rmax, name=None)

TODO: add doc.

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.tf.env.op_module.SoftMinVirial(du, sw_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.tf.env.op_module.TabulateFusion(table, table_info, em_x, em, last_layer_size, name=None)

TODO: add doc.

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.tf.env.op_module.TabulateFusionGrad(table, table_info, em_x, em, dy, descriptor, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.TabulateFusionGradGrad(table, table_info, em_x, em, dz_dy_dem_x, dz_dy_dem, descriptor, name=None)

TODO: add doc.

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.tf.env.op_module.TabulateFusionSeA(table, table_info, em_x, em, last_layer_size, name=None)

TODO: add doc.

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.tf.env.op_module.TabulateFusionSeAGrad(table, table_info, em_x, em, dy, descriptor, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.TabulateFusionSeAGradGrad(table, table_info, em_x, em, dz_dy_dem_x, dz_dy_dem, descriptor, is_sorted=True, name=None)

TODO: add doc.

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.

  • is_sorted – An optional bool. Defaults to True.

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

Returns

A Tensor. Has the same type as table.

deepmd.tf.env.op_module.TabulateFusionSeAtten(table, table_info, em_x, em, two_embed, last_layer_size, is_sorted=True, name=None)

TODO: add doc.

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.

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

  • last_layer_size – An int.

  • is_sorted – An optional bool. Defaults to True.

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

Returns

A Tensor. Has the same type as table.

deepmd.tf.env.op_module.TabulateFusionSeAttenGrad(table, table_info, em_x, em, two_embed, dy, descriptor, is_sorted=True, name=None)

TODO: add doc.

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.

  • two_embed – 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.

  • is_sorted – An optional bool. Defaults to True.

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

Returns

A tuple of Tensor objects (dy_dem_x, dy_dem, dy_dtwo).

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

deepmd.tf.env.op_module.TabulateFusionSeAttenGradGrad(table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem, dz_dy_dtwo, descriptor, is_sorted=True, name=None)

TODO: add doc.

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.

  • two_embed – 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.

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

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

  • is_sorted – An optional bool. Defaults to True.

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

Returns

A Tensor. Has the same type as table.

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

TODO: add doc.

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.tf.env.op_module.TabulateFusionSeRGrad(table, table_info, em, dy, descriptor, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.TabulateFusionSeRGradGrad(table, table_info, em, dz_dy_dem, descriptor, name=None)

TODO: add doc.

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.tf.env.op_module.TabulateFusionSeT(table, table_info, em_x, em, last_layer_size, name=None)

TODO: add doc.

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.tf.env.op_module.TabulateFusionSeTGrad(table, table_info, em_x, em, dy, descriptor, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.TabulateFusionSeTGradGrad(table, table_info, em_x, em, dz_dy_dem_x, dz_dy_dem, descriptor, name=None)

TODO: add doc.

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.tf.env.op_module.Tanh4FltNvnmd(x, name=None)

TODO: add doc.

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.tf.env.op_module.UnaggregatedDy2Dx(z, w, dy_dx, dy2_dx, ybar, functype, name=None)

TODO: add doc.

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.tf.env.op_module.UnaggregatedDy2DxS(y, dy, w, xbar, functype, name=None)

TODO: add doc.

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.tf.env.op_module.UnaggregatedDyDx(z, w, dy_dx, ybar, functype, name=None)

TODO: add doc.

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.tf.env.op_module.UnaggregatedDyDxS(y, w, xbar, functype, name=None)

TODO: add doc.

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.tf.env.op_module.add_flt_nvnmd(x: Any, w: Any, name=None) Any

TODO: add doc.

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

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

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.convert_forward_map(sub_forward_map: Any, sub_natoms: Any, natoms: Any, name=None)

TODO: add doc.

Parameters
  • sub_forward_map – A Tensor of type int32.

  • sub_natoms – A Tensor of type int32.

  • natoms – A Tensor of type int32.

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

Returns

A tuple of Tensor objects (forward_map, backward_map, new_natoms, mesh).

forward_map: A Tensor of type int32. backward_map: A Tensor of type int32. new_natoms: A Tensor of type int32. mesh: A Tensor of type int32.

deepmd.tf.env.op_module.copy_flt_nvnmd(x: Any, name=None)

TODO: add doc.

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

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

Returns

A tuple of Tensor objects (y1, y2).

y1: A Tensor. Has the same type as x. y2: A Tensor. Has the same type as x.

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

TODO: add doc.

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.tf.env.op_module.descrpt_norot(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, davg: Any, dstd: Any, rcut_a: float, rcut_r: float, rcut_r_smth: float, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.env.op_module.descrpt_se_a(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, davg: Any, dstd: Any, rcut_a: float, rcut_r: float, rcut_r_smth: float, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.env.op_module.descrpt_se_a_ef(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, ef: Any, davg: Any, dstd: Any, rcut_a: float, rcut_r: float, rcut_r_smth: float, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.env.op_module.descrpt_se_a_ef_para(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, ef: Any, davg: Any, dstd: Any, rcut_a: float, rcut_r: float, rcut_r_smth: float, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.env.op_module.descrpt_se_a_ef_vert(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, ef: Any, davg: Any, dstd: Any, rcut_a: float, rcut_r: float, rcut_r_smth: float, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.env.op_module.descrpt_se_a_mask(coord: Any, type: Any, mask: Any, box: Any, natoms: Any, mesh: Any, name=None)

TODO: add doc.

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

  • type – A Tensor of type int32.

  • mask – A Tensor of type int32.

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

  • natoms – A Tensor of type int32.

  • mesh – A Tensor of type int32.

  • 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.tf.env.op_module.descrpt_se_r(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, davg: Any, dstd: Any, rcut: float, rcut_smth: float, sel, name=None)

TODO: add doc.

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.tf.env.op_module.dotmul_flt_nvnmd(x: Any, w: Any, name=None) Any

TODO: add doc.

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

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

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.dprc_pairwise_idx(idxs: Any, natoms: Any, name=None)

TODO: add doc.

Parameters
  • idxs – A Tensor of type int32.

  • natoms – A Tensor of type int32.

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

Returns

A tuple of Tensor objects (forward_qm_map, backward_qm_map, forward_qmmm_map, backward_qmmm_map, natoms_qm, natoms_qmmm, qmmm_frame_idx).

forward_qm_map: A Tensor of type int32. backward_qm_map: A Tensor of type int32. forward_qmmm_map: A Tensor of type int32. backward_qmmm_map: A Tensor of type int32. natoms_qm: A Tensor of type int32. natoms_qmmm: A Tensor of type int32. qmmm_frame_idx: A Tensor of type int32.

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

TODO: add doc.

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.tf.env.op_module.flt_nvnmd(x: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.gelu(x: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.gelu_custom(x: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.gelu_grad(dy: Any, x: Any, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_module.gelu_grad_custom(dy: Any, x: Any, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_module.gelu_grad_grad(dy: Any, dy_: Any, x: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.gelu_grad_grad_custom(dy: Any, dy_: Any, x: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.map_aparam(aparam: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_module.map_flt_nvnmd(x: Any, table: Any, table_grad: Any, table_info: Any, name=None) Any

TODO: add doc.

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

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

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

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

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.matmul_fitnet_nvnmd(x: Any, w: Any, nbitx: int, nbitw: int, normw: int, name=None) Any

TODO: add doc.

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

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

  • nbitx – An int.

  • nbitw – An int.

  • normw – An int.

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.matmul_flt2fix_nvnmd(x: Any, w: Any, nbit: int, name=None) Any

TODO: add doc.

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

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

  • nbit – An int.

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.matmul_flt_nvnmd(x: Any, w: Any, normx: int, normw: int, name=None) Any

TODO: add doc.

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

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

  • normx – An int.

  • normw – An int.

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

Returns

A Tensor. Has the same type as x.

deepmd.tf.env.op_module.mul_flt_nvnmd(x: Any, w: Any, name=None) Any

TODO: add doc.

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

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

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

Returns

A Tensor. Has the same type as x.

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

TODO: add doc.

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.tf.env.op_module.pair_tab(table_info: Any, table_data: Any, type: Any, rij: Any, nlist: Any, natoms: Any, scale: Any, sel_a, sel_r, name=None)

TODO: add doc.

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.tf.env.op_module.parallel_prod_force_se_a(net_deriv: Any, in_deriv: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, parallel: bool = False, start_frac: float = 0, end_frac: float = 1, name=None) Any

TODO: add doc.

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.tf.env.op_module.prod_env_mat_a(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, davg: Any, dstd: Any, rcut_a: float, rcut_r: float, rcut_r_smth: float, 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.tf.env.op_module.prod_env_mat_a_mix(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, davg: Any, dstd: Any, rcut_a: float, rcut_r: float, rcut_r_smth: float, 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.tf.env.op_module.prod_env_mat_a_mix_nvnmd_quantize(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, davg: Any, dstd: Any, rcut_a: float, rcut_r: float, rcut_r_smth: float, sel_a, sel_r, name=None)

TODO: add doc.

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. nmask: A Tensor of type bool.

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

TODO: add doc.

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.tf.env.op_module.prod_env_mat_r(coord: Any, type: Any, natoms: Any, box: Any, mesh: Any, davg: Any, dstd: Any, rcut: float, rcut_smth: float, sel, name=None)

TODO: add doc.

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.tf.env.op_module.prod_force(net_deriv: Any, in_deriv: Any, nlist: Any, axis: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.prod_force_norot(net_deriv: Any, in_deriv: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.prod_force_se_a(net_deriv: Any, in_deriv: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.prod_force_se_a_mask(net_deriv: Any, in_deriv: Any, mask: Any, nlist: Any, total_atom_num: int, name=None) Any

TODO: add doc.

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.

  • mask – A Tensor of type int32.

  • nlist – A Tensor of type int32.

  • total_atom_num – An int.

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

Returns

A Tensor. Has the same type as net_deriv.

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

TODO: add doc.

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.tf.env.op_module.prod_virial(net_deriv: Any, in_deriv: Any, rij: Any, nlist: Any, axis: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None)

TODO: add doc.

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.tf.env.op_module.prod_virial_norot(net_deriv: Any, in_deriv: Any, rij: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None)

TODO: add doc.

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.tf.env.op_module.prod_virial_se_a(net_deriv: Any, in_deriv: Any, rij: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None)

TODO: add doc.

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.tf.env.op_module.prod_virial_se_r(net_deriv: Any, in_deriv: Any, rij: Any, nlist: Any, natoms: Any, name=None)

TODO: add doc.

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.tf.env.op_module.quantize_nvnmd(x: Any, isround: int, nbit1: int, nbit2: int, nbit3: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.soft_min_force(du: Any, sw_deriv: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.soft_min_switch(type: Any, rij: Any, nlist: Any, natoms: Any, sel_a, sel_r, alpha: float, rmin: float, rmax: float, name=None)

TODO: add doc.

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.tf.env.op_module.soft_min_virial(du: Any, sw_deriv: Any, rij: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None)

TODO: add doc.

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.tf.env.op_module.tabulate_fusion(table: Any, table_info: Any, em_x: Any, em: Any, last_layer_size: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.tabulate_fusion_grad(table: Any, table_info: Any, em_x: Any, em: Any, dy: Any, descriptor: Any, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.tabulate_fusion_grad_grad(table: Any, table_info: Any, em_x: Any, em: Any, dz_dy_dem_x: Any, dz_dy_dem: Any, descriptor: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.tabulate_fusion_se_a(table: Any, table_info: Any, em_x: Any, em: Any, last_layer_size: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.tabulate_fusion_se_a_grad(table: Any, table_info: Any, em_x: Any, em: Any, dy: Any, descriptor: Any, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.tabulate_fusion_se_a_grad_grad(table: Any, table_info: Any, em_x: Any, em: Any, dz_dy_dem_x: Any, dz_dy_dem: Any, descriptor: Any, is_sorted: bool = True, name=None) Any

TODO: add doc.

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.

  • is_sorted – An optional bool. Defaults to True.

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

Returns

A Tensor. Has the same type as table.

deepmd.tf.env.op_module.tabulate_fusion_se_atten(table: Any, table_info: Any, em_x: Any, em: Any, two_embed: Any, last_layer_size: int, is_sorted: bool = True, name=None) Any

TODO: add doc.

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.

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

  • last_layer_size – An int.

  • is_sorted – An optional bool. Defaults to True.

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

Returns

A Tensor. Has the same type as table.

deepmd.tf.env.op_module.tabulate_fusion_se_atten_grad(table: Any, table_info: Any, em_x: Any, em: Any, two_embed: Any, dy: Any, descriptor: Any, is_sorted: bool = True, name=None)

TODO: add doc.

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.

  • two_embed – 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.

  • is_sorted – An optional bool. Defaults to True.

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

Returns

A tuple of Tensor objects (dy_dem_x, dy_dem, dy_dtwo).

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

deepmd.tf.env.op_module.tabulate_fusion_se_atten_grad_grad(table: Any, table_info: Any, em_x: Any, em: Any, two_embed: Any, dz_dy_dem_x: Any, dz_dy_dem: Any, dz_dy_dtwo: Any, descriptor: Any, is_sorted: bool = True, name=None) Any

TODO: add doc.

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.

  • two_embed – 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.

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

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

  • is_sorted – An optional bool. Defaults to True.

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

Returns

A Tensor. Has the same type as table.

deepmd.tf.env.op_module.tabulate_fusion_se_r(table: Any, table_info: Any, em: Any, last_layer_size: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.tabulate_fusion_se_r_grad(table: Any, table_info: Any, em: Any, dy: Any, descriptor: Any, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_module.tabulate_fusion_se_r_grad_grad(table: Any, table_info: Any, em: Any, dz_dy_dem: Any, descriptor: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.tabulate_fusion_se_t(table: Any, table_info: Any, em_x: Any, em: Any, last_layer_size: int, name=None) Any

TODO: add doc.

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.tf.env.op_module.tabulate_fusion_se_t_grad(table: Any, table_info: Any, em_x: Any, em: Any, dy: Any, descriptor: Any, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_module.tabulate_fusion_se_t_grad_grad(table: Any, table_info: Any, em_x: Any, em: Any, dz_dy_dem_x: Any, dz_dy_dem: Any, descriptor: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.tanh4_flt_nvnmd(x: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.unaggregated_dy2_dx(z: Any, w: Any, dy_dx: Any, dy2_dx: Any, ybar: Any, functype: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.unaggregated_dy2_dx_s(y: Any, dy: Any, w: Any, xbar: Any, functype: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.unaggregated_dy_dx(z: Any, w: Any, dy_dx: Any, ybar: Any, functype: Any, name=None) Any

TODO: add doc.

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.tf.env.op_module.unaggregated_dy_dx_s(y: Any, w: Any, xbar: Any, functype: Any, name=None) Any

TODO: add doc.

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.

op_grads_module

Python wrappers around TensorFlow ops.

This file is MACHINE GENERATED! Do not edit.

class deepmd.tf.env.op_grads_module.Annotated(*args, **kwargs)

Bases: object

Add context specific metadata to a type.

Example: Annotated[int, runtime_check.Unsigned] indicates to the hypothetical runtime_check module that this type is an unsigned int. Every other consumer of this type can ignore this metadata and treat this type as int.

The first argument to Annotated must be a valid type.

Details:

  • It’s an error to call Annotated with less than two arguments.

  • Nested Annotated are flattened:

    Annotated[Annotated[T, Ann1, Ann2], Ann3] == Annotated[T, Ann1, Ann2, Ann3]
    
  • Instantiating an annotated type is equivalent to instantiating the

underlying type:

Annotated[C, Ann1](5) == C(5)
  • Annotated can be used as a generic type alias:

    Optimized = Annotated[T, runtime.Optimize()]
    Optimized[int] == Annotated[int, runtime.Optimize()]
    
    OptimizedList = Annotated[List[T], runtime.Optimize()]
    OptimizedList[int] == Annotated[List[int], runtime.Optimize()]
    
deepmd.tf.env.op_grads_module.ProdForceGrad(grad, net_deriv, in_deriv, nlist, axis, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.ProdForceSeAGrad(grad, net_deriv, in_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.ProdForceSeAMaskGrad(grad, net_deriv, in_deriv, mask, nlist, total_atom_num, name=None)

TODO: add doc.

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.

  • mask – A Tensor of type int32.

  • nlist – A Tensor of type int32.

  • total_atom_num – An int.

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

Returns

A Tensor. Has the same type as grad.

deepmd.tf.env.op_grads_module.ProdForceSeRGrad(grad, net_deriv, in_deriv, nlist, natoms, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.ProdVirialGrad(grad, net_deriv, in_deriv, rij, nlist, axis, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.ProdVirialSeAGrad(grad, net_deriv, in_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.ProdVirialSeRGrad(grad, net_deriv, in_deriv, rij, nlist, natoms, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.SoftMinForceGrad(grad, du, sw_deriv, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.SoftMinVirialGrad(grad, du, sw_deriv, rij, nlist, natoms, n_a_sel, n_r_sel, name=None)

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.prod_force_grad(grad: Any, net_deriv: Any, in_deriv: Any, nlist: Any, axis: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.prod_force_se_a_grad(grad: Any, net_deriv: Any, in_deriv: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.prod_force_se_a_mask_grad(grad: Any, net_deriv: Any, in_deriv: Any, mask: Any, nlist: Any, total_atom_num: int, name=None) Any

TODO: add doc.

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.

  • mask – A Tensor of type int32.

  • nlist – A Tensor of type int32.

  • total_atom_num – An int.

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

Returns

A Tensor. Has the same type as grad.

deepmd.tf.env.op_grads_module.prod_force_se_r_grad(grad: Any, net_deriv: Any, in_deriv: Any, nlist: Any, natoms: Any, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.prod_virial_grad(grad: Any, net_deriv: Any, in_deriv: Any, rij: Any, nlist: Any, axis: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.prod_virial_se_a_grad(grad: Any, net_deriv: Any, in_deriv: Any, rij: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.prod_virial_se_r_grad(grad: Any, net_deriv: Any, in_deriv: Any, rij: Any, nlist: Any, natoms: Any, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.soft_min_force_grad(grad: Any, du: Any, sw_deriv: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.

deepmd.tf.env.op_grads_module.soft_min_virial_grad(grad: Any, du: Any, sw_deriv: Any, rij: Any, nlist: Any, natoms: Any, n_a_sel: int, n_r_sel: int, name=None) Any

TODO: add doc.

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.