Python tensorflow.segment_sum() Examples

The following are 30 code examples of tensorflow.segment_sum(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module tensorflow , or try the search function .
Example #1
Source File: graph_layers.py    From PADME with MIT License 6 votes vote down vote up
def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """
    parent layers: atom_features, atom_membership
    """
    if in_layers is None:
      in_layers = self.in_layers
    in_layers = convert_to_layers(in_layers)

    self.build()
    output = in_layers[0].out_tensor
    atom_membership = in_layers[1].out_tensor
    for i, W in enumerate(self.W_list[:-1]):
      output = tf.matmul(output, W) + self.b_list[i]
      output = self.activation(output)
    output = tf.matmul(output, self.W_list[-1]) + self.b_list[-1]
    if self.output_activation:
      output = self.activation(output)
    output = tf.segment_sum(output, atom_membership)
    out_tensor = output
    if set_tensors:
      self.variables = self.trainable_weights
      self.out_tensor = out_tensor
    return out_tensor 
Example #2
Source File: graph_layers.py    From PADME with MIT License 6 votes vote down vote up
def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """
    parent layers: atom_features, membership
    """
    if in_layers is None:
      in_layers = self.in_layers
    in_layers = convert_to_layers(in_layers)

    # Add trainable weights
    self.build()

    # Extract atom_features
    atom_features = in_layers[0].out_tensor
    membership = in_layers[1].out_tensor
    # Extract atom_features
    graph_features = tf.segment_sum(atom_features, membership)
    # sum all graph outputs
    outputs = self.DAGgraph_step(graph_features, self.W_list, self.b_list)
    out_tensor = outputs
    if set_tensors:
      self.variables = self.trainable_weights
      self.out_tensor = out_tensor
    return out_tensor 
Example #3
Source File: layers.py    From graph_level_drug_discovery with Apache License 2.0 6 votes vote down vote up
def call(self, x):
    """Execute this layer on input tensors.

    Parameters
    ----------
    x: list of Tensor
      should be [embedding tensor of molecules, of shape (batch_size*max_n_atoms*n_embedding),
                 mask tensor of molecules, of shape (batch_size*max_n_atoms)]

    Returns
    -------
    list of tf.Tensor
      Of shape (batch_size)
    """
    self.build()
    output = x[0]
    atom_membership = x[1]
    for i, W in enumerate(self.W_list[:-1]):
      output = tf.matmul(output, W) + self.b_list[i]
      output = self.activation(output)
    output = tf.matmul(output, self.W_list[-1]) + self.b_list[-1]
    if self.output_activation:
      output = self.activation(output)
    output = tf.segment_sum(output, atom_membership)
    return output 
Example #4
Source File: layers.py    From graph_level_drug_discovery with Apache License 2.0 6 votes vote down vote up
def call(self, x, mask=None):
    """Execute this layer on input tensors.

    x = [graph_features, membership]

    Parameters
    ----------
    x: tf.Tensor
      Tensor of each atom's graph features

    Returns
    -------
    outputs: tf.Tensor
      Tensor of each molecule's features

    """
    # Add trainable weights
    self.build()
    atom_features = x[0]
    membership = x[1]
    # Extract atom_features
    graph_features = tf.segment_sum(atom_features, membership)
    # sum all graph outputs
    outputs = self.DAGgraph_step(graph_features, self.W_list, self.b_list)
    return outputs 
Example #5
Source File: layers_copy.py    From graph_level_drug_discovery with Apache License 2.0 6 votes vote down vote up
def call(self, x):
    """Execute this layer on input tensors.

    Parameters
    ----------
    x: list of Tensor
      should be [embedding tensor of molecules, of shape (batch_size*max_n_atoms*n_embedding),
                 mask tensor of molecules, of shape (batch_size*max_n_atoms)]

    Returns
    -------
    list of tf.Tensor
      Of shape (batch_size)
    """
    self.build()
    output = x[0]
    atom_membership = x[1]
    for i, W in enumerate(self.W_list[:-1]):
      output = tf.matmul(output, W) + self.b_list[i]
      output = self.activation(output)
    output = tf.matmul(output, self.W_list[-1]) + self.b_list[-1]
    if self.output_activation:
      output = self.activation(output)
    output = tf.segment_sum(output, atom_membership)
    return output 
Example #6
Source File: graph_layers.py    From PADME with MIT License 5 votes vote down vote up
def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """ 
    parent layers: atom_features, atom_split
    """
    if in_layers is None:
      in_layers = self.in_layers
    in_layers = convert_to_layers(in_layers)

    self.build()
    outputs = in_layers[0].out_tensor
    atom_split = in_layers[1].out_tensor

    if self.gaussian_expand:
      outputs = self.gaussian_histogram(outputs)

    output_molecules = tf.segment_sum(outputs, atom_split)

    if self.gaussian_expand:
      output_molecules = tf.matmul(output_molecules, self.W) + self.b
      output_molecules = self.activation(output_molecules)

    out_tensor = output_molecules
    if set_tensors:
      self.variables = self.trainable_weights
      self.out_tensor = out_tensor
    return out_tensor 
Example #7
Source File: graph_layers.py    From PADME with MIT License 5 votes vote down vote up
def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """
    parent layers: atom_features, distance, distance_membership_i, distance_membership_j
    """
    if in_layers is None:
      in_layers = self.in_layers
    in_layers = convert_to_layers(in_layers)

    self.build()
    atom_features = in_layers[0].out_tensor
    distance = in_layers[1].out_tensor
    distance_membership_i = in_layers[2].out_tensor
    distance_membership_j = in_layers[3].out_tensor
    distance_hidden = tf.matmul(distance, self.W_df) + self.b_df
    atom_features_hidden = tf.matmul(atom_features, self.W_cf) + self.b_cf
    outputs = tf.multiply(distance_hidden,
                          tf.gather(atom_features_hidden,
                                    distance_membership_j))

    # for atom i in a molecule m, this step multiplies together distance info of atom pair(i,j)
    # and embeddings of atom j(both gone through a hidden layer)
    outputs = tf.matmul(outputs, self.W_fc)
    outputs = self.activation(outputs)

    output_ii = tf.multiply(self.b_df, atom_features_hidden)
    output_ii = tf.matmul(output_ii, self.W_fc)
    output_ii = self.activation(output_ii)

    # for atom i, sum the influence from all other atom j in the molecule
    outputs = tf.segment_sum(outputs,
                             distance_membership_i) - output_ii + atom_features
    out_tensor = outputs
    if set_tensors:
      self.variables = self.trainable_weights
      self.out_tensor = out_tensor
    return out_tensor 
Example #8
Source File: graph_layers.py    From PADME with MIT License 5 votes vote down vote up
def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """ Perform M steps of set2set gather,
        detailed descriptions in: https://arxiv.org/abs/1511.06391 """
    if in_layers is None:
      in_layers = self.in_layers
    in_layers = convert_to_layers(in_layers)

    self.build()
    # Extract atom_features
    atom_features = in_layers[0].out_tensor
    atom_split = in_layers[1].out_tensor

    self.c = tf.zeros((self.batch_size, self.n_hidden))
    self.h = tf.zeros((self.batch_size, self.n_hidden))

    for i in range(self.M):
      q_expanded = tf.gather(self.h, atom_split)
      e = tf.reduce_sum(atom_features * q_expanded, 1)
      e_mols = tf.dynamic_partition(e, atom_split, self.batch_size)
      # Add another value(~-Inf) to prevent error in softmax
      e_mols = [
          tf.concat([e_mol, tf.constant([-1000.])], 0) for e_mol in e_mols
      ]
      a = tf.concat([tf.nn.softmax(e_mol)[:-1] for e_mol in e_mols], 0)
      r = tf.segment_sum(tf.reshape(a, [-1, 1]) * atom_features, atom_split)
      # Model using this layer must set pad_batches=True
      q_star = tf.concat([self.h, r], axis=1)
      self.h, self.c = self.LSTMStep(q_star, self.c)

    out_tensor = q_star
    if set_tensors:
      self.variables = self.trainable_weights
      self.out_tensor = out_tensor
    return out_tensor 
Example #9
Source File: ops.py    From tfdeploy with MIT License 5 votes vote down vote up
def test_SegmentSum(self):
        t = tf.segment_sum(self.random(4, 2, 3), np.array([0, 1, 1, 2]))
        self.check(t) 
Example #10
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testValues(self):
    dtypes = [tf.float32,
              tf.float64,
              tf.int64,
              tf.int32,
              tf.complex64,
              tf.complex128]

    # Each item is np_op1, np_op2, tf_op
    ops_list = [(np.add, None, tf.segment_sum),
                (self._mean_cum_op, self._mean_reduce_op,
                 tf.segment_mean),
                (np.ndarray.__mul__, None, tf.segment_prod),
                (np.minimum, None, tf.segment_min),
                (np.maximum, None, tf.segment_max)]

    # A subset of ops has been enabled for complex numbers
    complex_ops_list = [(np.add, None, tf.segment_sum),
                        (np.ndarray.__mul__, None, tf.segment_prod)]

    n = 10
    shape = [n, 2]
    indices = [i // 3 for i in range(n)]
    for dtype in dtypes:
      if dtype in (tf.complex64, tf.complex128):
        curr_ops_list = complex_ops_list
      else:
        curr_ops_list = ops_list

      with self.test_session(use_gpu=False):
        tf_x, np_x = self._input(shape, dtype=dtype)
        for np_op1, np_op2, tf_op in curr_ops_list:
          np_ans = self._segmentReduce(indices, np_x, np_op1, np_op2)
          s = tf_op(data=tf_x, segment_ids=indices)
          tf_ans = s.eval()
          self._assertAllClose(indices, np_ans, tf_ans)
          # NOTE(mrry): The static shape inference that computes
          # `tf_ans.shape` can only infer that sizes from dimension 1
          # onwards, because the size of dimension 0 is data-dependent
          # and may therefore vary dynamically.
          self.assertAllEqual(np_ans.shape[1:], tf_ans.shape[1:]) 
Example #11
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsShape(self):
    shape = [4, 4]
    tf_x, _ = self._input(shape)
    indices = tf.constant([0, 1, 2, 2], shape=[2, 2])
    with self.assertRaises(ValueError):
      tf.segment_sum(data=tf_x, segment_ids=indices) 
Example #12
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsSize(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 1]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment_ids should be the same size"):
        s.eval() 
Example #13
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsValid(self):
    # This is a baseline for the following SegmentIdsInvalid* tests.
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 0, 0, 1]
      result = tf.segment_sum(data=tf_x, segment_ids=indices).eval()
      self.assertAllEqual([[15, 18, 21, 24], [13, 14, 15, 16]], result) 
Example #14
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsInvalid1(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [-1, -1, 0, 0]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids do not start at 0"):
        s.eval() 
Example #15
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsInvalid3(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 0, 2, 2]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids are not increasing by 1"):
        s.eval() 
Example #16
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsInvalid4(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 1, 0, 1]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids are not increasing by 1"):
        s.eval() 
Example #17
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsInvalid5(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 1, 2, 0]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError(
          r"Segment id 1 out of range \[0, 1\), probably "
          "because 'segment_ids' input is not sorted."):
        s.eval() 
Example #18
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsInvalid6(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 0, 0, -1]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids must be >= 0"):
        s.eval() 
Example #19
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testSegmentIdsInvalid7(self):
    shape = [4, 4]
    with self.test_session():
      tf_x, _ = self._input(shape)
      indices = [0, 0, 0, -2]
      s = tf.segment_sum(data=tf_x, segment_ids=indices)
      with self.assertRaisesOpError("segment ids must be >= 0"):
        s.eval() 
Example #20
Source File: segment_reduction_ops_test.py    From deep_image_model with Apache License 2.0 5 votes vote down vote up
def testGradientMatchesSegmentSum(self):
    # Strategy: compute the gradient for UnsortedSegmentSum and SegmentSum
    # and compare the outputs, which should be identical.
    # NB: for this test to work, indices must be valid for SegmentSum, namely
    # it must be sorted, the indices must be contiguous, and num_segments
    # must be max(indices) + 1.
    indices = [0, 0, 1, 1, 1, 2, 3, 4, 5]
    n = len(indices)
    num_cols = 2
    shape = [n, num_cols]
    num_segments = max(indices) + 1
    with self.test_session(use_gpu=self.use_gpu):
      tf_x, np_x = self._input(shape, dtype=tf.float64)
      # Results from UnsortedSegmentSum
      unsorted_s = tf.unsorted_segment_sum(data=tf_x,
                                                 segment_ids=indices,
                                                 num_segments=num_segments)
      (unsorted_jacob_t, unsorted_jacob_n) = tf.test.compute_gradient(
          tf_x,
          shape,
          unsorted_s,
          [num_segments, num_cols],
          x_init_value=np_x.astype(np.double),
          delta=1)
      # Results from SegmentSum
      sorted_s = tf.segment_sum(data=tf_x, segment_ids=indices)
      sorted_jacob_t, sorted_jacob_n = tf.test.compute_gradient(
          tf_x,
          shape,
          sorted_s,
          [num_segments, num_cols],
          x_init_value=np_x.astype(np.double),
          delta=1)
    self.assertAllClose(unsorted_jacob_t, sorted_jacob_t, rtol=1e-3, atol=1e-3)
    self.assertAllClose(unsorted_jacob_n, sorted_jacob_n, rtol=1e-3, atol=1e-3) 
Example #21
Source File: models.py    From scene-graph-TF-release with MIT License 5 votes vote down vote up
def _compute_vert_context_soft(self, edge_factor, vert_factor, reuse=False):
        """
        attention-based vertex(node) message pooling
        """

        out_edge = utils.pad_and_gather(edge_factor, self.edge_pair_mask_inds[:,0])
        in_edge = utils.pad_and_gather(edge_factor, self.edge_pair_mask_inds[:,1])
        # gather correspounding vert factors
        vert_factor_gathered = tf.gather(vert_factor, self.edge_pair_segment_inds)

        # concat outgoing edges and ingoing edges with gathered vert_factors
        out_edge_w_input = tf.concat(concat_dim=1, values=[out_edge, vert_factor_gathered])
        in_edge_w_input = tf.concat(concat_dim=1, values=[in_edge, vert_factor_gathered])

        # compute compatibility scores
        (self.feed(out_edge_w_input)
             .fc(1, relu=False, reuse=reuse, name='out_edge_w_fc')
             .sigmoid(name='out_edge_score'))
        (self.feed(in_edge_w_input)
             .fc(1, relu=False, reuse=reuse, name='in_edge_w_fc')
             .sigmoid(name='in_edge_score'))

        out_edge_w = self.get_output('out_edge_score')
        in_edge_w = self.get_output('in_edge_score')

        # weight the edge factors with computed weigths
        out_edge_weighted = tf.mul(out_edge, out_edge_w)
        in_edge_weighted = tf.mul(in_edge, in_edge_w)


        edge_sum = out_edge_weighted + in_edge_weighted
        vert_ctx = tf.segment_sum(edge_sum, self.edge_pair_segment_inds)
        return vert_ctx 
Example #22
Source File: pooling.py    From SchNet with MIT License 5 votes vote down vote up
def __init__(self, mode='sum', name=None):
        if mode == 'sum':
            self._reduce = tf.segment_sum
        elif mode == 'mean':
            self._reduce = tf.segment_mean
        super(PoolSegments, self).__init__(name) 
Example #23
Source File: grimme_d3.py    From PhysNet with MIT License 5 votes vote down vote up
def _ncoord(Zi, Zj, r, idx_i, cutoff=None, k1=d3_k1, rcov=d3_rcov):
    '''
    compute coordination numbers by adding an inverse damping function
    '''
    rco = tf.gather(rcov,Zi) + tf.gather(rcov,Zj)
    rr = tf.cast(rco,r.dtype)/r
    damp = 1.0/(1.0+tf.exp(-k1*(rr-1.0)))
    if cutoff is not None:
        damp *= _smootherstep(r, cutoff)
    return tf.segment_sum(damp,idx_i) 
Example #24
Source File: NeuralNetwork.py    From PhysNet with MIT License 5 votes vote down vote up
def energy_from_scaled_atomic_properties(self, Ea, Qa, Dij, Z, idx_i, idx_j, batch_seg=None):
        with tf.name_scope("energy_from_atomic_properties"):
            if batch_seg is None:
                batch_seg = tf.zeros_like(Z)
            #add electrostatic and dispersion contribution to atomic energy
            if self.use_electrostatic:
                Ea += self.electrostatic_energy_per_atom(Dij, Qa, idx_i, idx_j)
            if self.use_dispersion:
                if self.lr_cut is not None:   
                    Ea += d3_autoev*edisp(Z, Dij/d3_autoang, idx_i, idx_j, s6=self.s6, s8=self.s8, a1=self.a1, a2=self.a2, cutoff=self.lr_cut/d3_autoang)
                else:
                    Ea += d3_autoev*edisp(Z, Dij/d3_autoang, idx_i, idx_j, s6=self.s6, s8=self.s8, a1=self.a1, a2=self.a2)
        return tf.squeeze(tf.segment_sum(Ea, batch_seg))

    #calculates the energy and forces given the scaled atomic atomic properties (in order to prevent recomputation if atomic properties are calculated) 
Example #25
Source File: NeuralNetwork.py    From PhysNet with MIT License 5 votes vote down vote up
def scaled_charges(self, Z, Qa, Q_tot=None, batch_seg=None):
        with tf.name_scope("scaled_charges"):
            if batch_seg is None:
                batch_seg = tf.zeros_like(Z)
            #number of atoms per batch (needed for charge scaling)
            Na_per_batch = tf.segment_sum(tf.ones_like(batch_seg, dtype=self.dtype), batch_seg)
            if Q_tot is None: #assume desired total charge zero if not given
                Q_tot = tf.zeros_like(Na_per_batch, dtype=self.dtype)
            #return scaled charges (such that they have the desired total charge)
            return Qa + tf.gather(((Q_tot-tf.segment_sum(Qa, batch_seg))/Na_per_batch), batch_seg)

    #switch function for electrostatic interaction (switches between shielded and unshielded electrostatic interaction) 
Example #26
Source File: NeuralNetwork.py    From PhysNet with MIT License 5 votes vote down vote up
def electrostatic_energy_per_atom(self, Dij, Qa, idx_i, idx_j):
        #gather charges
        Qi = tf.gather(Qa, idx_i)
        Qj = tf.gather(Qa, idx_j)
        #calculate variants of Dij which we need to calculate
        #the various shileded/non-shielded potentials
        DijS = tf.sqrt(Dij*Dij + 1.0) #shielded distance
        #calculate value of switching function
        switch = self._switch(Dij) #normal switch
        cswitch = 1.0-switch #complementary switch
        #calculate shielded/non-shielded potentials
        if self.lr_cut is None: #no non-bonded cutoff
            Eele_ordinary = 1.0/Dij   #ordinary electrostatic energy
            Eele_shielded = 1.0/DijS  #shielded electrostatic energy
            #combine shielded and ordinary interactions and apply prefactors 
            Eele = self.kehalf*Qi*Qj*(cswitch*Eele_shielded + switch*Eele_ordinary)
        else: #with non-bonded cutoff
            cut   = self.lr_cut
            cut2  = self.lr_cut*self.lr_cut
            Eele_ordinary = 1.0/Dij  +  Dij/cut2 - 2.0/cut
            Eele_shielded = 1.0/DijS + DijS/cut2 - 2.0/cut
            #combine shielded and ordinary interactions and apply prefactors 
            Eele = self.kehalf*Qi*Qj*(cswitch*Eele_shielded + switch*Eele_ordinary)
            Eele = tf.where(Dij <= cut, Eele, tf.zeros_like(Eele))
        return tf.segment_sum(Eele, idx_i) 

    #save the current model 
Example #27
Source File: utils.py    From PiNN with BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def get_atomic_dress(dataset, elems, max_iter=None):
    """Fit the atomic energy with a element dependent atomic dress

    Args:
        dataset: dataset to fit
        elems: a list of element numbers
    Returns:
        atomic_dress: a dictionary comprising the atomic energy of each element
        error: residue error of the atomic dress
    """
    tensors = dataset.make_one_shot_iterator().get_next()
    if 'ind_1' not in tensors:
        tensors['ind_1'] = tf.expand_dims(tf.zeros_like(tensors['elems']), 1)
        tensors['e_data'] = tf.expand_dims(tensors['e_data'], 0)
    count = tf.equal(tf.expand_dims(
        tensors['elems'], 1), tf.expand_dims(elems, 0))
    count = tf.cast(count, tf.int32)
    count = tf.segment_sum(count, tensors['ind_1'][:, 0])
    sess = tf.Session()
    x, y = [], []
    it = 0
    while True:
        it += 1
        if max_iter is not None and it > max_iter:
            break
        try:
            x_i, y_i = sess.run((count, tensors['e_data']))
            x.append(x_i)
            y.append(y_i)
        except tf.errors.OutOfRangeError:
            break
    x, y = np.concatenate(x, 0), np.concatenate(y, 0)
    beta = np.dot(np.dot(np.linalg.pinv(np.dot(x.T, x)), x.T), np.array(y))
    dress = {e: float(beta[i]) for (i, e) in enumerate(elems)}
    error = np.dot(x, beta) - y
    return dress, error 
Example #28
Source File: ops.py    From document-qa with Apache License 2.0 5 votes vote down vote up
def segment_logsumexp(xs, segments):
    """ Similar tf.segment_sum but compute logsumexp rather then sum """
    # Stop gradients following the implementation of tf.reduce_logsumexp
    maxs = tf.stop_gradient(tf.reduce_max(xs, axis=1))
    segment_maxes = tf.segment_max(maxs, segments)
    xs -= tf.expand_dims(tf.gather(segment_maxes, segments), 1)
    sums = tf.reduce_sum(tf.exp(xs), axis=1)
    return tf.log(tf.segment_sum(sums, segments)) + segment_maxes 
Example #29
Source File: graph_layers.py    From PADME with MIT License 4 votes vote down vote up
def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):
    """ Creates weave tensors.

    parent layers: [atom_features, pair_features], pair_split, atom_to_pair
    """
    activation = activations.get(self.activation)
    if in_layers is None:
      in_layers = self.in_layers
    in_layers = convert_to_layers(in_layers)

    self.build()

    atom_features = in_layers[0].out_tensor
    pair_features = in_layers[1].out_tensor

    pair_split = in_layers[2].out_tensor
    atom_to_pair = in_layers[3].out_tensor

    AA = tf.matmul(atom_features, self.W_AA) + self.b_AA
    AA = activation(AA)
    PA = tf.matmul(pair_features, self.W_PA) + self.b_PA
    PA = activation(PA)
    PA = tf.segment_sum(PA, pair_split)

    A = tf.matmul(tf.concat([AA, PA], 1), self.W_A) + self.b_A
    A = activation(A)

    if self.update_pair:
      AP_ij = tf.matmul(
          tf.reshape(
              tf.gather(atom_features, atom_to_pair),
              [-1, 2 * self.n_atom_input_feat]), self.W_AP) + self.b_AP
      AP_ij = activation(AP_ij)
      AP_ji = tf.matmul(
          tf.reshape(
              tf.gather(atom_features, tf.reverse(atom_to_pair, [1])),
              [-1, 2 * self.n_atom_input_feat]), self.W_AP) + self.b_AP
      AP_ji = activation(AP_ji)

      PP = tf.matmul(pair_features, self.W_PP) + self.b_PP
      PP = activation(PP)
      P = tf.matmul(tf.concat([AP_ij + AP_ji, PP], 1), self.W_P) + self.b_P
      P = activation(P)
    else:
      P = pair_features

    self.out_tensors = [A, P]
    #pdb.set_trace()
    if set_tensors:
      self.variables = self.trainable_weights
      self.out_tensor = A
    return self.out_tensors 
Example #30
Source File: segments.py    From moonlight with Apache License 2.0 4 votes vote down vote up
def _segments_1d(values, mode, name=None):
  """Labels consecutive runs of the same value.

  Args:
    values: 1D tensor of any type.
    mode: The SegmentsMode. Returns the start of each segment (STARTS), or the
      rounded center of each segment (CENTERS).
    name: Optional name for the op.

  Returns:
    run_centers: int32 tensor; the centers of each run with the same consecutive
        values.
    run_lengths: int32 tensor; the lengths of each run.

  Raises:
    ValueError: if mode is not recognized.
  """
  with tf.name_scope(name, "segments", [values]):

    def do_segments(values):
      """Actually does segmentation.

      Args:
        values: 1D tensor of any type. Non-empty.

      Returns:
        run_centers: int32 tensor
        run_lengths: int32 tensor

      Raises:
        ValueError: if mode is not recognized.
      """
      length = tf.shape(values)[0]
      values = tf.convert_to_tensor(values)
      # The first run has id 0, so we don't increment the id.
      # Otherwise, the id is incremented when the value changes.
      run_start_bool = tf.concat(
          [[False], tf.not_equal(values[1:], values[:-1])], axis=0)
      # Cumulative sum the run starts to get the run ids.
      segment_ids = tf.cumsum(tf.cast(run_start_bool, tf.int32))
      if mode is SegmentsMode.STARTS:
        run_centers = tf.segment_min(tf.range(length), segment_ids)
      elif mode is SegmentsMode.CENTERS:
        run_centers = tf.segment_mean(
            tf.cast(tf.range(length), tf.float32), segment_ids)
        run_centers = tf.cast(tf.floor(run_centers), tf.int32)
      else:
        raise ValueError("Unexpected mode: %s" % mode)
      run_lengths = tf.segment_sum(tf.ones([length], tf.int32), segment_ids)
      return run_centers, run_lengths

    def empty_segments():
      return (tf.zeros([0], tf.int32), tf.zeros([0], tf.int32))

    return tf.cond(
        tf.greater(tf.shape(values)[0], 0), lambda: do_segments(values),
        empty_segments)