def build_step(self, signals):
input = signals.gather(self.input_data)
input = tf.reshape(input, (self.n_ops, -1))
state = signals.gather(self.state_sig)
# compute output
if self.C is None:
output = tf.zeros_like(input)
else:
output = state * self.C
output = tf.reshape(
output,
(self.n_ops, -1, signals.minibatch_size * self.signal_d))
output = tf.reduce_sum(output, axis=1)
if self.D is not None:
output += self.D * input
signals.scatter(self.output_data, output)
# update state
r = gen_sparse_ops._sparse_tensor_dense_mat_mul(
self.A_indices, self.A, self.A_shape, state)
with tf.control_dependencies([output]):
state = r + tf.scatter_nd(self.offsets, input,
self.state_sig.shape)
# TODO: tensorflow does not yet support sparse_tensor_dense_add
# on the GPU
# state = gen_sparse_ops._sparse_tensor_dense_add(
# self.offsets, input, self.state_sig.shape, r)
state.set_shape(self.state_sig.shape)
signals.mark_gather(self.input_data)
signals.mark_gather(self.state_sig)
signals.scatter(self.state_sig, state)
评论列表
文章目录