def _ct_ht_shared_variable_finder(self, context):
if self.lstm_cell_type == RNNUnitType.LSTMBlockCell:
return None
lstm_cell = context.cell_match
ct = lstm_cell.get_op("ct").output[0]
ht = lstm_cell.get_op("ht").output[0]
ct_concat = [
c for c in self.g.find_output_consumers(ct) if is_concat_op(c)
]
ht_concat = [
c for c in self.g.find_output_consumers(ht) if is_concat_op(c)
]
if len(ct_concat) != 1 or len(
ht_concat) != 1 or ct_concat[0] != ht_concat[0]:
log.debug("failed to find ct-ht concat")
return None
ct_ht_shared_output = ct_concat[0].output[0]
consumers = []
ct_identity_consumer = lstm_cell.get_op("ct_identity_consumer")
ht_identity_consumer = lstm_cell.get_op("xh")
ct_slice = [c for c in ct_identity_consumer.inputs if is_slice_op(c)]
ht_slice = [c for c in ht_identity_consumer.inputs if is_slice_op(c)]
if len(ct_slice) != 1 or len(ht_slice) != 1:
log.debug("failed to find slice op before identity consumers")
return None
consumers.extend([ct_slice[0], ht_slice[0]])
return self._find_state_variable_with_select(context,
ct_ht_shared_output,
consumers)
def find_inputs(self, rnn_scope_name, rnn_props, match, input_blacklist=None):
concat_node = match.get_op("cell_inputs")
assert is_concat_op(concat_node)
read_node = concat_node.inputs[0]
assert is_tensor_array_read_op(read_node)
enter_node = read_node.inputs[2]
assert enter_node.type == "Enter"
scatter_node = enter_node.inputs[0]
assert is_tensor_array_scatter_op(scatter_node)
node = scatter_node.inputs[2]
node_id = scatter_node.input[2]
# dynamic_rnn may insert transpose op if input data format is [B, T, D]
if node.type == "Transpose" and node.name.startswith(rnn_scope_name):
node_id = node.input[0]
node = node.inputs[0]
assert not node.name.startswith(rnn_scope_name)
rnn_props.input_node = node
rnn_props.input_id = node_id