class_docstr: Deep Tensor Neural Networks
  
  This class implements deep tensor neural networks as first defined in [1]_
  
  References
  ----------
  .. [1] Schütt, Kristof T., et al. "Quantum-chemical insights from deep
      tensor neural networks." Nature communications 8.1 (2017): 1-8.
class_docstr_tok: ["Deep","Tensor","Neural","Networks","This","class","implements","deep","tensor","neural","networks","as","first","defined","in","[","1","]","_","References","--","--","--","--","--","..","[","1","]","Sch\u00fctt",",","Kristof","T.",",","et","al",".","``","Quantum-chemical","insights","from","deep","tensor","neural","networks",".","''","Nature","communications","8.1","(","2017",")",":","1-8","."]
docstr: Computes the values for different Feature Layers on given batch
  
  A tf.py_func wrapper is written around this when creating the
  input_fn for tf.Estimator
docstr_tok: ["Computes","the","values","for","different","Feature","Layers","on","given","batch","A","tf.py_func","wrapper","is","written","around","this","when","creating","the","input_fn","for","tf.Estimator"]
code: def compute_features_on_batch(self, X_b):
      """Computes the values for different Feature Layers on given batch
  
      A tf.py_func wrapper is written around this when creating the
      input_fn for tf.Estimator
  
      """
      distance = []
      atom_membership = []
      distance_membership_i = []
      distance_membership_j = []
      num_atoms = list(map(sum, X_b.astype(bool)[:, :, 0]))
      atom_number = [
          np.round(
              np.power(
                  2 * np.diag(X_b[i, : num_atoms[i], : num_atoms[i]]),
                  1 / 2.4,
              )
          ).astype(int)
          for i in range(len(num_atoms))
      ]
      start = 0
      for im, molecule in enumerate(atom_number):
          distance_matrix = (
              np.outer(molecule, molecule)
              / X_b[im, : num_atoms[im], : num_atoms[im]]
          )
          np.fill_diagonal(distance_matrix, -100)
          distance.append(np.expand_dims(distance_matrix.flatten(), 1))
          atom_membership.append([im] * num_atoms[im])
          membership = np.array(
              [np.arange(num_atoms[im])] * num_atoms[im]
          )
          membership_i = membership.flatten(order="F")
          membership_j = membership.flatten()
          distance_membership_i.append(membership_i + start)
          distance_membership_j.append(membership_j + start)
          start = start + num_atoms[im]
  
      atom_number = np.concatenate(atom_number).astype(np.int32)
      distance = np.concatenate(distance, axis=0)
      gaussian_dist = np.exp(
          -np.square(distance - self.steps) / (2 * self.step_size**2)
      )
      gaussian_dist = gaussian_dist.astype(np.float32)
      atom_mem = np.concatenate(atom_membership).astype(np.int32)
      dist_mem_i = np.concatenate(distance_membership_i).astype(
          np.int32
      )
      dist_mem_j = np.concatenate(distance_membership_j).astype(
          np.int32
      )
  
      features = [
          atom_number,
          gaussian_dist,
          atom_mem,
          dist_mem_i,
          dist_mem_j,
      ]
  
      return features
  
imports: from collections.abc import Sequence
  import deepchem
  import numpy
  import tensorflow
  from typing import List, Union, Tuple, Iterable, Dict, Optional
  from deepchem.utils.typing import OneOrMany, LossFn, ActivationFn
  from deepchem.data import Dataset, pad_features
  from deepchem.feat.mol_graphs import ConvMol
  from deepchem.metrics import to_one_hot
  from deepchem.models import KerasModel, layers
  from deepchem.models.losses import L2Loss, SoftmaxCrossEntropy, Loss
  from tensorflow.keras.layers import Input, Dense, Reshape, Softmax, Dropout, Activation, BatchNormalization
  import warnings
label_desc: Use image node_id 2 for calling the DTNNModel obj's underlying member method code with example usage: obj.compute_features_on_batch(X_b) and returns: features