code: def DistanceTensor(X, Nbrs, boxsize, B, N, M, d):
      """Calculates distance tensor for batch of molecules.
  
      B = batch_size, N = max_num_atoms, M = max_num_neighbors, d = num_features
  
      Parameters
      ----------
      X: tf.Tensor of shape (B, N, d)
        Coordinates/features tensor.
      Nbrs: tf.Tensor of shape (B, N, M)
        Neighbor list tensor.
      boxsize: float or None
        Simulation box length [Angstrom].
  
      Returns
      -------
      D: tf.Tensor of shape (B, N, M, d)
        Coordinates/features distance tensor.
  
      """
  
      atom_tensors = tf.unstack(X, axis=1)
      nbr_tensors = tf.unstack(Nbrs, axis=1)
      D = []
      if boxsize is not None:
          for atom, atom_tensor in enumerate(atom_tensors):
              nbrs = gather_neighbors(X, nbr_tensors[atom], B, N, M, d)
              nbrs_tensors = tf.unstack(nbrs, axis=1)
              for nbr, nbr_tensor in enumerate(nbrs_tensors):
                  _D = tf.subtract(nbr_tensor, atom_tensor)
                  _D = tf.subtract(
                      _D,
                      boxsize * tf.round(tf.math.divide(_D, boxsize)),
                  )
                  D.append(_D)
      else:
          for atom, atom_tensor in enumerate(atom_tensors):
              nbrs = gather_neighbors(X, nbr_tensors[atom], B, N, M, d)
              nbrs_tensors = tf.unstack(nbrs, axis=1)
              for nbr, nbr_tensor in enumerate(nbrs_tensors):
                  _D = tf.subtract(nbr_tensor, atom_tensor)
                  D.append(_D)
      D = tf.stack(D)
      D = tf.transpose(D, perm=[1, 0, 2])
      D = tf.reshape(D, [B, N, M, d])
      return D
  
file_path: deepchem/contrib/atomicconv/models/atomicnet_ops.py
imports: from __future__ import print_function
  from __future__ import division
  from __future__ import unicode_literals
  import numpy
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: DistanceTensor(X, Nbrs, boxsize, B, N, M, d) and returns: D