yapf attempt

  cells: tf.Tensor

    (box_size**ndim, ndim) shape.

  """

  return tf.reshape(

      tf.transpose(

          tf.stack(

              tf.meshgrid(

                  *[tf.range(start, stop, nbr_cutoff) for _ in range(ndim)]))),

      (-1, ndim))

  ranges = [tf.range(start, stop, nbr_cutoff) for _ in range(ndim)]

  return tf.reshape(tf.transpose(tf.stack(tf.meshgrid(*ranges))), (-1, ndim))

def put_atoms_in_cells(coords, cells, N, n_cells, ndim, k=5):

def hydrophobic(d):

  """Compute hydrophobic interaction term."""

  return tf.where(d < 0.5, tf.ones_like(d),

                  tf.where(d < 1.5, 1.5 - d, tf.zeros_like(d)))

  where = tf.where(d < 1.5, 1.5 - d, tf.zeros_like(d))

  return tf.where(d < 0.5, tf.ones_like(d), where)

def hbond(d):

  """Computes hydrogen bond term."""

  return tf.where(d < -0.7, tf.ones_like(d),

                  tf.where(d < 0, (1.0 / 0.7) * (0 - d), tf.zeros_like(d)))

  where = tf.where(d < 0, (1.0 / 0.7) * (0 - d), tf.zeros_like(d))

  return tf.where(d < -0.7, tf.ones_like(d), where)

def g(c, Nrot):

      gc1 = GraphConv(

          layer_size,

          activation_fn=tf.nn.relu,

          in_layers=[in_layer, self.degree_slice, self.membership] +

          self.deg_adjs)

          in_layers=[in_layer, self.degree_slice, self.membership

                    ] + self.deg_adjs)

      batch_norm1 = BatchNorm(in_layers=[gc1])

      in_layer = GraphPool(

          in_layers=[batch_norm1, self.degree_slice, self.membership] +