a437991e12a334b40b082fdc8b1b6dcf4a68021f,deepchem/models/tensorgraph/layers.py,NeighborList,get_neighbor_cells,#NeighborList#Any#,1021

Before Change

    // Lists of length n_cells
    coords_rel = [
        tf.to_float(cells) - tf.to_float(centers)
        for (cells, centers) in zip(tiled_centers, tiled_cells)
    ]
    coords_norm = [tf.reduce_sum(rel**2, axis=1) for rel in coords_rel]

    // Lists of length n_cells
    // Get indices of n_nbrs atoms closest to each cell point
    // n_cells tensors of shape (n_nbrs,)
    closest_inds = tf.stack([tf.nn.top_k(norm, k=n_nbrs)[1] for norm in coords_norm])

    return closest_inds

After Change

    // Tile (a, b, c, a, b, c, ...)
    tiled_cells = tf.tile(cells, (n_cells, 1))

    coords_vec = tf.reduce_sum((tiled_centers - tiled_cells)**2, axis=1)
    coords_norm = tf.reshape(coords_vec, (n_cells, n_cells))
    closest_inds = tf.nn.top_k(-coords_norm, k=n_nbr_cells)[1]

    return closest_inds

Instances