def call( self, inputs: Tuple[Union[tf.Tensor, List[tf.Tensor], tf.SparseTensor]], ) -> tf.Tensor: """Returns called Graph Convolution Layer. Parameters --------------------------- inputs: Tuple[Union[tf.Tensor, tf.SparseTensor]], """ adjacency, node_features = inputs[0], inputs[1:] ids = tf.SparseTensor(indices=adjacency.indices, values=adjacency.indices[:, 1], dense_shape=adjacency.dense_shape) return [ self._l2_norm( dense( embedding_ops.embedding_lookup_sparse_v2( self._dropout_layer(node_feature), ids, adjacency, combiner='mean'))) for dense, node_feature in zip(self._dense_layers, node_features) ]
def call(self, inputs): if inputs.dtype.base_dtype != self._compute_dtype_object.base_dtype: inputs = math_ops.cast(inputs, dtype=self._compute_dtype_object) rank = inputs.shape.rank if rank == 2 or rank is None: # We use embedding_lookup_sparse as a more efficient matmul operation for # large sparse input tensors. The op will result in a sparse gradient, as # opposed to sparse_ops.sparse_tensor_dense_matmul which results in dense # gradients. This can lead to sigfinicant speedups, see b/171762937. if isinstance(inputs, sparse_tensor.SparseTensor): # We need to fill empty rows, as the op assumes at least one id per row. inputs, _ = sparse_ops.sparse_fill_empty_rows(inputs, 0) # We need to do some munging of our input to use the embedding lookup as # a matrix multiply. We split our input matrix into separate ids and # weights tensors. The values of the ids tensor should be the column # indices of our input matrix and the values of the weights tensor # can continue to the actual matrix weights. # The column arrangement of ids and weights # will be summed over and does not matter. See the documentation for # sparse_ops.sparse_tensor_dense_matmul a more detailed explanation # of the inputs to both ops. ids = sparse_tensor.SparseTensor( indices=inputs.indices, values=inputs.indices[:, 1], dense_shape=inputs.dense_shape) weights = inputs outputs = embedding_ops.embedding_lookup_sparse_v2( self.kernel * self.window, ids, weights, combiner='sum') else: outputs = gen_math_ops.MatMul(a=inputs, b=self.kernel * self.window) # Broadcast kernel to inputs. else: outputs = standard_ops.tensordot(inputs, self.kernel * self.window, [[rank - 1], [0]]) # Reshape the output back to the original ndim of the input. if not context.executing_eagerly(): shape = inputs.shape.as_list() output_shape = shape[:-1] + [self.kernel.shape[-1]] outputs.set_shape(output_shape) if self.use_bias: outputs = nn_ops.bias_add(outputs, self.bias) if self.activation is not None: outputs = self.activation(outputs) return outputs
def _embedding_lookup_for_sparse_tensor( inp: sparse_tensor.SparseTensor, weight: Optional[sparse_tensor.SparseTensor], table: tf_variables.Variable, feature: tpu_embedding_v2_utils.FeatureConfig) -> ops.Tensor: """Embedding lookup for sparse tensor based on its feature config. Args: inp: a single SparseTensor input. weight: None or SparseTensor which has the same shape of the input. table: a table variable. feature: a feature config. Returns: Embedding lookup result. """ if not feature.output_shape and feature.max_sequence_length > 0: batch_size = math_ops.cast(array_ops.shape(inp)[0], dtype=dtypes.int64) sparse_shape = array_ops.stack( [batch_size, feature.max_sequence_length], axis=0) # TPU Embedding truncates sequences to max_sequence_length, and if we # don't truncate, scatter_nd will error out if the index was out of # bounds. truncated_inp = sparse_ops.sparse_slice(inp, start=[0, 0], size=sparse_shape) dense_output_shape = array_ops.stack( [batch_size, feature.max_sequence_length, feature.table.dim], axis=0) return array_ops.scatter_nd( truncated_inp.indices, array_ops.gather(table.read_value(), truncated_inp.values), dense_output_shape) else: inp_rank = inp.dense_shape.get_shape()[0] if (not feature.validate_weights_and_indices and inp_rank is not None and inp_rank <= 2): return embedding_ops.embedding_lookup_sparse_v2( table, inp, sp_weights=weight, combiner=feature.table.combiner) else: return embedding_ops.safe_embedding_lookup_sparse_v2( table, inp, sparse_weights=weight, combiner=feature.table.combiner)
def sparse_lookup(): sp_ids = sparse_tensor.SparseTensor( indices=[[0, 0], [0, 1], [1, 0], [2, 2]], values=[0, 3, 4, 1], dense_shape=[3, 3]) return embedding_ops.embedding_lookup_sparse_v2(sv, sp_ids, None)
def dense(inputs, kernel, bias=None, activation=None, dtype=None): """Densely connected NN layer op. Args: inputs: `tf.Tensor` or `tf.SparseTensor`. Inputs to operation. kernel: `tf.Variable`. Matrix kernel. bias: (Optional) `tf.Variable`. Bias to add to outputs. activation: (Optional) 1-argument callable. Activation function to apply to outputs. dtype: (Optional) `tf.DType`. Dtype to cast `inputs` to. Returns: `tf.Tensor`. Output of dense connection. """ if dtype: if inputs.dtype.base_dtype != dtype.base_dtype: inputs = math_ops.cast(inputs, dtype=dtype) rank = inputs.shape.rank if rank == 2 or rank is None: # We use embedding_lookup_sparse as a more efficient matmul operation for # large sparse input tensors. The op will result in a sparse gradient, as # opposed to sparse_ops.sparse_tensor_dense_matmul which results in dense # gradients. This can lead to sigfinicant speedups, see b/171762937. if isinstance(inputs, sparse_tensor.SparseTensor): # We need to fill empty rows, as the op assumes at least one id per row. inputs, _ = sparse_ops.sparse_fill_empty_rows(inputs, 0) # We need to do some munging of our input to use the embedding lookup as a # matrix multiply. We split our input matrix into separate ids and weights # tensors. The values of the ids tensor should be the column indices of # our input matrix and the values of the weights tensor can continue to # the actual matrix weights. The column arrangement of ids and weights # will be summed over and does not matter. See the documentation for # sparse_ops.sparse_tensor_dense_matmul a more detailed explanation of the # inputs to both ops. ids = sparse_tensor.SparseTensor(indices=inputs.indices, values=inputs.indices[:, 1], dense_shape=inputs.dense_shape) weights = inputs outputs = embedding_ops.embedding_lookup_sparse_v2(kernel, ids, weights, combiner="sum") else: outputs = gen_math_ops.MatMul(a=inputs, b=kernel) # Broadcast kernel to inputs. else: outputs = standard_ops.tensordot(inputs, kernel, [[rank - 1], [0]]) # Reshape the output back to the original ndim of the input. if not context.executing_eagerly(): shape = inputs.shape.as_list() output_shape = shape[:-1] + [kernel.shape[-1]] outputs.set_shape(output_shape) if bias is not None: outputs = nn_ops.bias_add(outputs, bias) if activation is not None: outputs = activation(outputs) return outputs